Регистрация операторов

Шаг 1

Упомянутые выше операторы имеют общую особенность: это встроенные операторы PyTorch с определёнными namespaces и Schemas, реализации которых уже существуют для встроенных ускорителей PyTorch (CPU, CUDA и т. д.). Далее нам нужно реализовать эти операторы для новых ускорителей.

 1at::Tensor empty_memory_format(
 2    c10::IntArrayRef size,
 3    std::optional<c10::ScalarType> dtype_opt,
 4    std::optional<c10::Layout> layout_opt,
 5    std::optional<c10::Device> device_opt,
 6    std::optional<bool> pin_memory_opt,
 7    std::optional<c10::MemoryFormat> memory_format_opt) {
 8  const auto device = c10::device_or_default(device_opt);
 9  const auto dtype = c10::dtype_or_default(dtype_opt);
10  TORCH_CHECK(device.is_privateuseone());
11  TORCH_CHECK(
12      c10::layout_or_default(layout_opt) == c10::Layout::Strided,
13      "Non strided layout not supported");
14  TORCH_CHECK(
15      !c10::pinned_memory_or_default(pin_memory_opt),
16      "Pin memory can only be on CPU");
17  const c10::DeviceGuard device_guard(device);
18  constexpr c10::DispatchKeySet pu1_dks(c10::DispatchKey::PrivateUse1);
19  auto allocator = at::GetAllocator(at::kPrivateUse1);
20  return at::detail::empty_generic(
21      size, allocator, pu1_dks, dtype, memory_format_opt);
22}

 1at::Tensor wrapper_empty_memory_format(
 2    c10::IntArrayRef size,
 3    std::optional<c10::ScalarType> dtype_opt,
 4    std::optional<c10::Layout> layout_opt,
 5    std::optional<c10::Device> device_opt,
 6    std::optional<bool> pin_memory_opt,
 7    std::optional<c10::MemoryFormat> memory_format_opt) {
 8  return at::native::openreg::empty_memory_format(
 9      size,
10      dtype_opt,
11      layout_opt,
12      device_opt,
13      pin_memory_opt,
14      memory_format_opt);
15}

В качестве примера возьмём оператор empty.memory_format. Сначала нужно запросить сведения schema об операторе в native_functions.yaml, где содержится подробная информация о сигнатуре. Затем можно реализовать оператор с учётом возможностей нового ускорителя.

- func: empty.memory_format(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor
dispatch:
    CPU: empty_cpu
    CUDA: empty_cuda
    ...

 1TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
 2  m.impl("empty.memory_format", wrapper_empty_memory_format);
 3  m.impl("empty_strided", wrapper_empty_strided);
 4  m.impl("as_strided", wrapper_as_strided);
 5  m.impl("resize_", wrapper_resize_);
 6  m.impl("_reshape_alias", wrapper__reshape_alias);
 7  m.impl("_copy_from", wrapper__copy_from);
 8  m.impl("_copy_from_and_resize", wrapper__copy_from_and_resize);
 9  m.impl("_local_scalar_dense", wrapper__local_scalar_densor);
10  m.impl(
11      "_has_compatible_shallow_copy_type",
12      wrapper_has_compatible_shallow_copy_type);
13  m.impl("set_.source_Tensor", wrapper_set_source_Tensor_);
14  m.impl("set_.source_Storage", wrapper_set_source_Storage_);
15  m.impl(
16      "set_.source_Storage_storage_offset",
17      wrapper_set_source_Storage_storage_offsetset_);
18  m.impl("view", wrapper_view);
19}

После завершения wrapper_empty_memory_format можно зарегистрировать aten::empty.memory_format для PrivateUse1 с помощью TORCH_LIBRARY_IMPL.

Шаг 2

Следуя инструкциям в разделе Шаг 1, можно завершить разработку и регистрацию всех операторов, кроме fallback. Далее, чтобы поддержать операторы, связанные с вычислениями (например, математические операции и свёртки), нужно реализовать регистрацию резервной семантики. Это встроенная возможность PyTorch, позволяющая передавать на выполнение на CPU некоторые операции, не поддерживаемые новыми ускорителями. Для новых разрабатываемых бэкендов это чрезвычайно эффективный способ обеспечить функциональность в ущерб производительности.

 1void cpu_fallback(const c10::OperatorHandle& op, torch::jit::Stack* stack) {
 2  static const std::unordered_set<c10::OperatorName> cpu_fallback_blocklist = {
 3      c10::OperatorName("aten::abs", ""),
 4      c10::OperatorName("aten::abs", "out"),
 5  };
 6
 7  const auto& op_name = op.schema().operator_name();
 8  if (cpu_fallback_blocklist.count(op_name)) {
 9    TORCH_CHECK(
10        false,
11        "Operator '",
12        op_name,
13        "' is not implemented for device openreg.");
14  } else {
15    at::native::cpu_fallback(op, stack);
16  }
17}

1void wrapper_cpu_fallback(
2    const c10::OperatorHandle& op,
3    torch::jit::Stack* stack) {
4  at::native::openreg::cpu_fallback(op, stack);
5}

1TORCH_LIBRARY_IMPL(_, PrivateUse1, m) {
2  m.fallback(
3      torch::CppFunction::makeFromBoxedFunction<&wrapper_cpu_fallback>());
4}

wrapper_cpu_fallback оборачивает метод at::native::cpu_fallback, предоставляемый PyTorch, и регистрируется в PyTorch с помощью PrivateUse1 через TORCH_LIBRARY_IMPL. Последующие операции, не поддерживаемые новым бэкендом, будут автоматически передаваться на выполнение CPU, а после выполнения результаты будут возвращаться новому бэкенду.

Выборочный резервный механизм

Вполне разумный сценарий — включить резервный механизм только для некоторых операторов, сохранив стандартное поведение PyTorch для остальных (если у ускорителя нет соответствующей реализации оператора, будет выдана ошибка).

1void wrapper_cpu_fallback(
2    const c10::OperatorHandle& op,
3    torch::jit::Stack* stack) {
4  at::native::openreg::cpu_fallback(op, stack);
5}

1TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
2  m.impl(
3      "sub.Tensor",
4      torch::CppFunction::makeFromBoxedFunction<&wrapper_cpu_fallback>());
5}

Резервные реализации для отдельных операторов очень похожи на глобальные резервные реализации; единственное различие — способ регистрации: вызов m.impl регистрирует реализацию для конкретного оператора, а m.fallback регистрирует реализацию по умолчанию для всех операторов.

 1void cpu_fallback(const c10::OperatorHandle& op, torch::jit::Stack* stack) {
 2  static const std::unordered_set<c10::OperatorName> cpu_fallback_blocklist = {
 3      c10::OperatorName("aten::abs", ""),
 4      c10::OperatorName("aten::abs", "out"),
 5  };
 6
 7  const auto& op_name = op.schema().operator_name();
 8  if (cpu_fallback_blocklist.count(op_name)) {
 9    TORCH_CHECK(
10        false,
11        "Operator '",
12        op_name,
13        "' is not implemented for device openreg.");
14  } else {
15    at::native::cpu_fallback(op, stack);
16  }
17}

Разумеется, глобальные резервные реализации можно также сочетать с их чёрным списком. Это распространённый подход, особенно если резервный механизм не поддерживается лишь несколькими операторами.

STUB в PyTorch

Для встроенных операторов PyTorch также предлагает другой подход: STUB. По сути, этот метод основан на подходе Шаг 1, но добавляет возможности вторичного планирования (например, планирование с учётом характеристик CPU).

pushd ${TORCH_ROOT}

find aten -type f -a -name "*.h" | xargs -I {} grep -wl "^DECLARE_DISPATCH" {}

popd

DECLARE_DISPATCH — это макрос, используемый для явного объявления STUB. Сейчас он находится в каталоге aten. С помощью этого макроса можно найти все операторы, которые можно интегрировать методом STUB.

...
aten/src/ATen/native/Activation.h
aten/src/ATen/native/FusedSGD.h
aten/src/ATen/native/nested/NestedTensorBinaryOps.h
aten/src/ATen/native/TensorCompare.h
aten/src/ATen/native/Sorting.h
...

using unary_fn = void(*)(TensorIteratorBase&);

DECLARE_DISPATCH(unary_fn, abs_stub)

В приведённом выше списке указан файл, в котором объявлен оператор STUB; в нём можно увидеть имя STUB и соответствующую сигнатуру функции. Далее на примере abs_stub мы кратко рассмотрим, как добавить поддержку операторов с помощью STUB.

 1void abs_kernel(at::TensorIteratorBase& iter) {
 2  TORCH_CHECK(iter.ntensors() == 2, "Abs kernel expects 2 tensors");
 3  TORCH_CHECK(
 4      iter.common_dtype() == at::ScalarType::Float,
 5      "Abs kernel only supports float type");
 6
 7  auto& output_tensor = iter.tensor(0);
 8  auto& input_tensor = iter.tensor(1);
 9
10  TORCH_CHECK(
11      input_tensor.sizes() == output_tensor.sizes(),
12      "Input and output tensor sizes must match.");
13
14  auto abs_loop = [](float* out_ptr, const float* in_ptr, int64_t n) {
15    for (int64_t i = 0; i < n; ++i) {
16      out_ptr[i] = std::abs(in_ptr[i]);
17    }
18  };
19
20  MemoryGuard guard(input_tensor, output_tensor);
21
22  if (iter.is_contiguous()) {
23    abs_loop(
24        static_cast<float*>(iter.data_ptr(0)),
25        static_cast<float*>(iter.data_ptr(1)),
26        iter.numel());
27  } else {
28    TORCH_CHECK(
29        input_tensor.is_contiguous(), "Input tensor must be contiguous.")
30
31    auto output = at::empty(
32        input_tensor.sizes(),
33        input_tensor.options().memory_format(
34            input_tensor.suggest_memory_format()));
35
36    MemoryGuard guard(output);
37
38    abs_loop(
39        static_cast<float*>(output.data_ptr()),
40        static_cast<float*>(iter.data_ptr(1)),
41        iter.numel());
42
43    output_tensor.copy_(output);
44  }
45}

1REGISTER_PRIVATEUSE1_DISPATCH(abs_stub, &wrapper_abs_stub);
2REGISTER_PRIVATEUSE1_DISPATCH(
3    quantize_tensor_per_tensor_affine_stub,
4    &wrapper_quantize_tensor_per_tensor_affine_stub);
5REGISTER_PRIVATEUSE1_DISPATCH(
6    _fused_sdp_choice_stub,
7    &wrapper__fused_sdp_choice);

Из сигнатуры видно, что входным параметром abs_stub является TensorIteratorBase — мощный вспомогательный класс PyTorch, содержащий все входные и выходные операторы, а также некоторые другие вспомогательные методы. На его основе можно разработать оператор abs_kernel, а затем вызвать REGISTER_PRIVATEUSE1_DISPATCH, чтобы указать abs_stub и завершить регистрацию.

Пользовательские операторы

Помимо встроенных операторов PyTorch, для повышения производительности в определённых сценариях часто используются пользовательские операторы ускорителей. Их можно разделить на три основных типа:

• Только прямой проход
• Прямой и обратный проходы: отдельная регистрация
• Прямой и обратный проходы: реализация с использованием torch.autograd.Function

Команды

PyTorch предоставляет набор команд с префиксом torch._C._dispatch_ для работы с функцией Dispatch. Запросить все связанные интерфейсы можно следующей командой:

python -c 'import torch; print("\n".join([x for x in dir(torch._C) if x.startswith("_dispatch_")]))'

...
_dispatch_dump
_dispatch_dump_table
_dispatch_has_kernel
_dispatch_has_kernel_for_any_dispatch_key
_dispatch_has_kernel_for_dispatch_key
_dispatch_isTensorSubclassLike
_dispatch_is_alias_key
_dispatch_is_included_in_alias
_dispatch_is_main_interpreter
_dispatch_kernel_for_dispatch_key_is_fallthrough
_dispatch_key_for_device
_dispatch_key_name
_dispatch_key_parse
_dispatch_key_set
...

Ниже приведены пояснения к нескольким часто используемым командам:

• torch._C._dispatch_key_set:

  Отображает DispatchKey текущего тензора; приоритет возрастает слева направо.

  >>> import torch
  >>> a = torch.randn(3,3,device="cuda")
  >>> torch._C._dispatch_key_set(a)
  'DispatchKeySet(CUDA, ADInplaceOrView, AutogradCUDA, AutocastCUDA)'
  

• torch._C._dispatch_dump_table:

  Проверяет поддержку заданного оператора для разных Dispatch Key, помогая найти соответствующий код реализации.

  >>> import torch
  >>> print(torch._C._dispatch_dump_table("aten::add.Tensor"))
  >>> ...
      CPU: registered at ./build/aten/src/ATen/RegisterCPU_0.cpp:1309 [kernel]
      CUDA: registered at ./build/aten/src/ATen/RegisterCUDA_0.cpp:2420 [kernel]
      HIP: registered at ./build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional_0.cpp:1373 [default backend kernel]
      MPS: registered at ./build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional_0.cpp:1373 [default backend kernel]
      IPU: registered at ./build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional_0.cpp:1373 [default backend kernel]
      XPU: registered at ./build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional_0.cpp:1373 [default backend kernel]
      HPU: registered at ./build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional_0.cpp:1373 [default backend kernel]
      VE: registered at ./build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional_0.cpp:1373 [default backend kernel]
      MTIA: registered at ./build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional_0.cpp:1373 [default backend kernel]
      MAIA: registered at ./build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional_0.cpp:1373 [default backend kernel]
      PrivateUse1: registered at ./build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional_0.cpp:1373 [default backend kernel]
      ...
  

  Вы можете легко запросить соответствующую реализацию оператора aten::add.Tensor на других платформах, чтобы проследить весь процесс вызова оператора на уровне исходного кода.

Переменные среды

В PyTorch также предусмотрены некоторые переменные среды, связанные с диспетчером, которые помогают изучать его работу и быстро находить проблемы.

• TORCH_SHOW_DISPATCH_TRACE

  Отображает подробные сведения о внутреннем планировании ключей диспетчеризации во время выполнения PyTorch.

  export TORCH_SHOW_DISPATCH_TRACE=1
  

  >>> import torch
  >>> a = torch.randn(3,3)
   [call] op=[aten::randn], key=[BackendSelect]
     [redispatch] op=[aten::randn], key=[CPU]
       [call] op=[aten::empty.memory_format], key=[BackendSelect]
         [redispatch] op=[aten::empty.memory_format], key=[CPU]
       [call] op=[aten::normal_], key=[CPU]
  

  Вы можете чётко увидеть все базовые операторы, вызываемые операторами уровня Python в PyTorch: в том числе имя оператора, иерархию вызовов и соответствующий Dispatch Key.