이 페이지는 Cloud Translation API를 통해 번역되었습니다.
Switch to English

XLA 사용자 지정 호출

이 문서는 XLA "사용자 정의 호출"을 작성하고 사용하는 방법을 설명합니다. 사용자 지정 호출을 사용하면 XLA 프로그램에서 C ++ 또는 CUDA와 같은 프로그래밍 언어로 작성된 코드를 호출 할 수 있습니다.

CPU 사용자 정의 호출

XLA의 클라이언트 API를 통해 사용자 지정 호출을 나타내는 HLO 명령어를 만들 수 있습니다. 이것은 글을 쓰는 시점에 TensorFlow를 통해 노출되지 않습니다.

예를 들어 다음 코드는 사용자 지정 호출을 사용하여 CPU에서 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 함수는 그것이 작동하는 버퍼의 크기를 알아야합니다. 이 예에서는 크기를 128 및 2048로 하드 코딩합니다.이 작업을 원하지 않으면 차원 매개 변수를 호출에 전달할 수 있습니다.

GPU의 맞춤형 호출

GPU 사용자 지정 호출 프레임 워크는 CPU의 프레임 워크와 약간 다릅니다. 다음은 위의 CPU 코드와 동일한 A[i] = B[i % 128] + C[i] 계산을 수행하는 CUDA 예제입니다.

 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");
 

먼저 GPU 사용자 정의 호출 기능 은 여전히 ​​CPU에서 실행되는 기능 입니다. do_custom_call CPU 함수는 GPU에서 작업을 대기열에 포함시킵니다. 여기서 CUDA 커널을 시작하지만 콜 큐브와 같은 다른 작업을 수행 할 수도 있습니다.

buffers 는 호스트에 존재하는 포인터의 배열이며, 각 요소는 포인트 (장치) (예 : GPU) 메모리를 포함합니다. 매개 변수가 먼저 나오고 그 뒤에 출력 값이옵니다. 이것은 insout 이라는 두 개의 매개 변수가있는 CPU 호출 규칙과 현저히 다릅니다. 우리가 분기하는 주된 이유는 튜플 모양의 입력 / 출력을 효율적으로 처리 할 수 ​​있기 때문입니다. 아래 섹션을 참조하십시오.

CPU 예제에서와 같이 입력 및 출력 버퍼 크기를 사용자 지정 호출에 하드 코딩했습니다. 그러나 CPU의 경우와 달리 피연산자로 버퍼 크기를 사용자 지정 호출에 전달하면 제대로 작동하지 않습니다. 일반적으로 CPU에서 사용할 수있는 버퍼 크기가 필요합니다. 예를 들어 커널을 시작할 때 사용할 블록 / 그리드 치수를 알아야합니다. 그러나 버퍼 크기를 피연산자로 사용자 지정 호출에 전달하면 값이 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 log를 확인하여 과거에 어떻게 변경되었는지 확인하십시오.

맞춤 호출에 튜플 전달

다음 사용자 정의 호출을 고려하십시오.

 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 op)이 여러 개의 독립 배열을 반환하는 방식입니다.

그러나 덜 분명하게도 튜플 출력은 사용자 정의 호출 임시 메모리를 제공하는 방법입니다. 예, 출력 은 임시 버퍼를 나타낼 수 있습니다. 출력 버퍼에는 op가 쓸 수있는 속성이 있으며, 작성된 후에는 읽을 수 있습니다. 그것이 바로 임시 버퍼에서 원하는 것입니다.

위의 예에서 F32[1024] 를 임시 버퍼로 사용하려고한다고 가정하십시오. 그런 다음 위와 같이 HLO를 작성하고 사용자 지정 호출 출력의 튜플 인덱스 1을 읽지 않았습니다.

CPU 사용자 지정 호출의 튜플

CPU 코드에는 do_custom_call(const void** ins, void* out) 함수가 있습니다. ins 는 하나의 요소를 가진 배열이며 param0 을 가리 킵니다. param0 의 서브 버퍼는 해당 포인터를 역 참조 out 액세스 할 수 있으며 output_tuple 의 서브 버퍼는 out 을 참조 out 액세스 할 수 있습니다.

GPU 사용자 지정 호출의 튜플

GPU 코드에는 do_custom_call(..., void** buffers, ...) 함수가 있습니다. 이 경우 buffers 는 입력 / 출력의 각 리프 버퍼마다 하나씩 6 개의 장치 포인터로 구성된 호스트 배열입니다. 단순 목록을 생성하려면 매개 변수와 출력을 반복하고 각각에 대해 셰이프의 사전 순회를 수행합니다. 구체적으로 :

 // 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