Эта страница была переведа с помощью Cloud Translation API.
Switch to English

Пользовательские звонки XLA

Этот документ описывает, как написать и использовать XLA «пользовательские вызовы». Пользовательские вызовы позволяют вам вызывать код, написанный на языке программирования, таком как C ++ или CUDA, из программы XLA.

Пользовательский вызов на CPU

Вы можете создать инструкцию HLO, которая представляет пользовательский вызов через клиентский API XLA. Это не раскрывается через TensorFlow на момент написания.

Например, следующий код использует пользовательский вызов для вычисления A[i] = B[i % 128] + C[i] на процессоре. (Конечно, вы можете - и должны! - делать это с помощью обычной HLO.)

 #include "tensorflow/compiler/xla/client/xla_builder.h"
#include "tensorflow/compiler/xla/service/custom_call_target_registry.h"

void do_it() {
  xla::XlaBuilder b("do_it");
  xla::XlaOp param0 =
      xla::Parameter(0, xla::ShapeUtil::CreateShape(F32, {128}), "p0");
  xla::XlaOp param1 =
      xla::Parameter(1, xla::ShapeUtil::CreateShape(F32, {2048}), "p1");
  xla::XlaOp custom_call =
      xla::CustomCall(&b, "do_custom_call", /*operands=*/{param0, param1},
                      /*output_shape=*/ShapeUtil::CreateShape(F32, {2048}));
}

void do_custom_call(void* out, const void** in) {
  float* out_buf = reinterpret_cast<float*>(out);
  const float* in0 = reinterpret_cast<const float*>(in[0]);
  const float* in1 = reinterpret_cast<const float*>(in[1]);
  for (int i = 0; i < 2048; ++i) {
    out_buf[i] = in0[i % 128] + in1[i];
  }
}
XLA_REGISTER_CUSTOM_CALL_TARGET(do_custom_call, "Host");
 

Обратите внимание, что функция do_custom_call должна знать размеры буферов, над do_custom_call она работает. В этом примере мы жестко закодируем размеры 128 и 2048. Если вы не хотите этого делать, вы можете передать измерения в качестве параметров для вызова.

Custom-вызов на GPU

Инфраструктура пользовательских вызовов GPU несколько отличается от инфраструктуры ЦП. Вот пример CUDA, который выполняет те же вычисления A[i] = B[i % 128] + C[i] что и код процессора выше.

 void do_it() { /* same implementation as above */ }

__global__ custom_call_kernel(const float* in0, const float* in1, float* out) {
  size_t idx = threadIdx.x * blockSize.x + gridIdx.x;
  out[idx] = in0[idx % 128] + in1[idx];
}

void do_custom_call(CUstream stream, void** buffers,
                    const char* opaque, size_t opaque_len) {
  const float* in0 = reinterpret_cast<const float*>(buffers[0]);
  const float* in1 = reinterpret_cast<const float*>(buffers[1]);
  float* out = reinterpret_cast<float*>(buffers[2]);

  const int64 block_dim = 64;
  const int64 grid_dim = 2048 / block_dim;
  custom_call_kernel<<<grid_dim, block_dim,
                       /*dynamic_shared_mem_bytes=*/0, stream>>>(in0, in1, out);
}
XLA_REGISTER_CUSTOM_CALL_TARGET(do_custom_call, "CUDA");
 

Сначала обратите внимание, что пользовательская функция вызова графического процессора все еще является функцией, выполняемой в ЦПУ . Наша do_custom_call процессора do_custom_call отвечает за постановку работы на GPU. Здесь он запускает ядро ​​CUDA, но может также делать что-то еще, например, вызывать cublas.

buffers - это массив указателей, которые живут на хосте, и каждый элемент, который он содержит, указывает на память устройства (т.е. GPU). Сначала идут параметры, а затем выходное значение. Это заметно отличается от соглашения о вызовах ЦП, которое имеет два параметра: ins и out . Основная причина, по которой мы расходимся, - это возможность эффективно обрабатывать входы / выходы в форме кортежей; см. раздел ниже.

Как и в примере с процессором, мы жестко закодировали размеры входного и выходного буферов в нашем пользовательском вызове. Однако, в отличие от случая с процессором, передача размеров буфера в качестве операндов в пользовательский вызов не будет работать хорошо. Обычно нам нужны размеры буфера, доступные нам на процессоре; Например, при запуске ядра нам нужно знать размеры блока / сетки, которые нужно использовать. Но если бы мы передали размеры буфера в качестве операндов нашему пользовательскому вызову, их значения остались бы в памяти GPU. Затем мы должны были бы сделать дорогой синхронный memcpy от устройства к хосту в начале нашей операции, чтобы просто прочитать размеры.

Чтобы вы могли обойти это, мы предоставляем opaque параметр. Вы можете установить это в произвольную строку байтов при создании пользовательского вызова:

 std::string opaque = "...";
xla::CustomCall(&b, "do_custom_call", /*operands=*/{param0, param1},
                /*output_shape=*/ShapeUtil::CreateShape(F32, {2048}),
                opaque);
 

Поскольку xla::Shape имеет представление буфера протокола, вы можете хранить этот сериализованный прото внутри opaque и десериализовать его в вашем пользовательском вызове GPU. Однако следует отметить , что , хотя xla::ShapeProto не меняется часто, это делает изменение. Проверьте журнал git, чтобы увидеть, как он изменился в прошлом.

Передача кортежей на пользовательские вызовы

Рассмотрим следующий пользовательский вызов.

 using xla::ShapeUtil;
Shape p0_shape = ShapeUtil::MakeTuple({
    ShapeUtil::MakeShape(F32, {32}),
    ShapeUtil::MakeTuple({
        ShapeUtil::MakeShape(F32, {64}),
        ShapeUtil::MakeShape(F32, {128}),
    }),
    ShapeUtil::MakeShape(F32, {256}),
});
xla::XlaOp p0 = xla::Parameter(0, p0_shape, "p0");

Shape out_shape = ShapeUtil::MakeTuple({
  ShapeUtil::MakeShape(F32, {512}),
  ShapeUtil::MakeShape(F32, {1024}),
});
xla::CustomCall(&b, "do_custom_call", /*operands=*/{p0}, out_shape);
 

И на CPU, и на GPU кортеж представлен в памяти в виде массива указателей. В C ++ - псевдокоде параметр 0 выше выложен следующим образом.

 // In-memory layout of parameter 0 from custom-call above.  True on both CPU
// and GPU.
float* subbuf0 = new float[32];
float* subbuf1 = new float[64];
float* subbuf2 = new float[128]
float* subbuf3 = new float[256];

void* subtuple = new void*[2];
(*subtuple)[0] = subbuf1;
(*subtuple)[1] = subbuf2;

void* p0 = new void*[3];
(*p0)[0] = subbuf0;
(*p0)[1] = subtuple;
(*p0)[2] = subbuf3;
 

Хотя представления кортежей в памяти одинаковы в CPU и GPU, они обрабатываются по-разному в соглашениях о вызовах пользовательских вызовов CPU и GPU.

Кортеж выводит как временные буферы

Кортежные входы для пользовательских вызовов удобны, но они не являются строго необходимыми. Если мы не поддерживаем ввод кортежей для пользовательских вызовов, вы всегда можете распаковать кортежи, используя get-tuple-element, прежде чем передавать их в пользовательский вызов.

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

Очевидная причина получения кортежей состоит в том, что пользовательский вызов (или любой другой оператор XLA) возвращает несколько независимых массивов.

Но менее очевидно, что вывод кортежей - это еще и способ предоставить пользовательскую временную память вызовов. Да, вывод может представлять временный буфер. Учтите, что у выходного буфера есть свойство, которое операционная система может записывать в него, и оно может читать из него после того, как оно было записано. Это именно то, что вы хотите от временного буфера.

В приведенном выше примере предположим, что мы хотели использовать F32[1024] в качестве временного буфера. Затем мы напишем HLO, как описано выше, и мы просто никогда не будем читать индекс кортежа 1 вывода пользовательского вызова.

Кортежи в процессорных вызовах

В коде процессора у нас есть функция do_custom_call(const void** ins, void* out) . ins - это массив только с одним элементом, который указывает на param0 . Подбуферы param0 доступны путем разыменования этого указателя, а подбуферы output_tuple доступны путем разыменования out .

Кортежи в графических процессорах

В коде GPU у нас есть функция do_custom_call(..., void** buffers, ...) . В этом случае buffers - это хост-массив из шести указателей устройства, по одному для каждого конечного буфера на входе / выходе. Чтобы сгенерировать плоский список, мы перебираем параметры и выходные данные, и для каждого делаем предварительный обход его формы. В частности:

 // Layout of `buffers` parameter to GPU custom call function for custom-call
// above.
buffers[0] == subbuf0
buffers[1] == subbuf1
buffers[2] == subbuf2
buffers[3] == subbuf3
buffers[4] == output_subbuf0
buffers[5] == output_subbuf1