XLA 맞춤 호출

이 문서에서는 XLA FFI 라이브러리를 사용하여 XLA 커스텀 호출을 작성하고 사용하는 방법을 설명합니다. 커스텀 호출은 컴파일 시간에 XLA 컴파일러에 HLO 모듈의 외부 '작업'을 설명하는 메커니즘이며, XLA FFI는 런타임에 XLA에 이러한 작업의 구현을 등록하는 메커니즘입니다. FFI는 '외부 함수 인터페이스'의 약자로, XLA가 다른 프로그래밍 언어로 작성된 외부 코드를 호출하기 위한 바이너리 인터페이스 (ABI)를 정의하는 C API 집합입니다. XLA는 C++로 작성된 XLA FFI의 헤더 전용 결합을 제공하여 최종 사용자로부터 기본 C API의 모든 낮은 수준의 세부정보를 숨깁니다.

JAX + XLA 커스텀 호출

JAX와 커스텀 호출 및 XLA FFI를 통합하는 엔드 투 엔드 예시는 JAX 문서를 참고하세요.

XLA FFI 결합

XLA FFI 결합은 커스텀 호출 서명(커스텀 호출 인수, 속성 및 유형, 실행 컨텍스트를 통해 전달되는 추가 매개변수(예: GPU 백엔드의 GPU 스트림))의 컴파일 시간 사양입니다. XLA FFI 결합은 호환되는 operator() 서명이 있는 모든 C++ 호출 가능 항목 (함수 포인터, 람다 등)에 결합될 수 있습니다. 생성된 핸들러는 안정적인 C API로 정의된 XLA FFI 호출 프레임을 디코딩하고, 모든 매개변수의 유형을 확인하고, 디코딩된 결과를 사용자 정의 콜백으로 전달합니다.

XLA FFI 결합은 생성된 핸들러를 가장 효율적인 기계어 코드로 컴파일할 수 있도록 템플릿 메타프로그래밍에 크게 의존합니다. 런타임 오버헤드는 각 커스텀 호출 매개변수에 대해 몇 나노초 정도입니다.

XLA FFI 맞춤설정 포인트는 템플릿 특수화로 구현되며, 사용자는 커스텀 유형을 디코딩하는 방법을 정의할 수 있습니다. 즉, 사용자 정의 enum class 유형에 대해 커스텀 디코딩을 정의할 수 있습니다.

커스텀 호출에서 오류 반환

커스텀 호출 구현은 XLA 런타임에 성공 또는 오류를 알리기 위해 xla::ffi::Error 값을 반환해야 합니다. 이는 absl::Status와 유사하며 동일한 오류 코드 집합을 포함합니다. absl::Status는 안정적인 ABI가 없으며 동적으로 로드된 커스텀 호출 라이브러리와 XLA 자체 간에 전달하는 것이 안전하지 않으므로 사용하지 않습니다.

// Handler that always returns an error.
auto always_error = Ffi::Bind().To(
    []() { return Error(ErrorCode::kInternal, "Oops!"); });

// Handler that always returns a success.
auto always_success = Ffi::Bind().To(
    []() { return Error::Success(); });

버퍼 인수 및 결과

XLA는 결과에 대상 전달 스타일을 사용합니다. 커스텀 호출 (또는 다른 XLA 작업)은 결과에 메모리를 할당하지 않고 대신 XLA 런타임에서 전달된 대상에 씁니다. XLA는 정적 버퍼 할당을 사용하며 컴파일 시간에 라이브 범위를 기반으로 모든 값에 버퍼를 할당합니다.

FFI 핸들러에 전달된 결과는 포인터와 유사한 시맨틱스를 갖는 Result<T> 템플릿으로 래핑됩니다. operator->는 기본 매개변수에 대한 액세스 권한을 제공합니다.

AnyBuffer 인수 및 결과는 모든 데이터 유형의 커스텀 호출 버퍼 매개변수에 대한 액세스 권한을 제공합니다. 이는 커스텀 호출에 여러 데이터 유형에서 작동하는 일반 구현이 있고 커스텀 호출 구현이 데이터 유형을 기반으로 런타임 디스패치를 실행하는 경우에 유용합니다. AnyBuffer 는 버퍼 데이터 유형, 측정기준, 버퍼 자체에 대한 포인터에 대한 액세스 권한을 제공합니다.

%0 = "stablehlo.custom_call"(%arg0) {
  call_target_name = "foo",
  api_version = 4 : i32
} : (tensor<2x2xf32>) -> tensor<2x2xf32>

// Buffers of any number of dimensions and data type.
auto handler = Ffi::Bind().Arg<AnyBuffer>().Ret<AnyBuffer>().To(
    [](AnyBuffer arg, Result<AnyBuffer> res) -> Error {
      void* arg_data = arg.untyped_data();
      void* res_data = res->untyped_data();
      return Error::Success();
    });

제약조건이 있는 버퍼 인수 및 결과

Buffer 를 사용하면 버퍼 데이터 유형 및 측정기준 수에 제약조건을 추가할 수 있으며, 런타임 인수가 FFI 핸들러 서명과 일치하지 않으면 핸들러에서 자동으로 확인하고 XLA 런타임에 오류를 반환합니다.

// Buffers of any number of dimensions and F32 data type.
auto handler = Ffi::Bind().Arg<Buffer<F32>>().Ret<Buffer<F32>>().To(
    [](Buffer<F32> arg, Result<Buffer<F32>> res) -> Error {
      float* arg_data = arg.typed_data();
      float* res_data = res->typed_data();
      return Error::Success();
    });

// Buffers of number of dimensions 2 and F32 data type.
auto handler = Ffi::Bind().Arg<BufferR2<F32>>().Ret<BufferR2<F32>>().To(
    [](BufferR2<F32> arg, Result<BufferR2<F32>> res) -> Error {
      float* arg_data = arg.typed_data();
      float* res_data = res->typed_data();
      return Error::Success();
    });

버퍼 일치 및 확인

버퍼 패턴은 FFI 결합의 인수 및 결과 유형보다 더 구체적인 제약조건을 표현할 수 있습니다. 이는 AnyBuffer를 구체적인 버퍼 유형으로 세분화하고, 측정기준 크기를 확인하고, 여러 버퍼 도형 간의 관계를 확인하는 데 유용합니다.

패턴은 변경 불가능하며 dtype 및 순위를 직접 또는 상응하는 수정자를 사용하여 지정할 수 있습니다.

namespace m = ::xla::ffi::match;

m::Buffer<F32, 2>();
m::Buffer().WithDType<F32>().WithRank<2>();

Match 함수는 패턴을 확인하고 최대한 제약조건이 있는 버퍼 유형을 반환합니다. AnyBuffer를 세분화하려면 패턴이 정확히 하나의 데이터 유형과 하나의 순위를 지정해야 합니다. 이미 유형이 지정된 버퍼와 일치하면 추가 제약조건이 확인되고 유형이 보존됩니다.

namespace m = ::xla::ffi::match;

auto handler = Ffi::Bind().Arg<AnyBuffer>().Ret<AnyBuffer>().To(
    [](AnyBuffer input, Result<AnyBuffer> output) -> Error {
      int64_t rows;
      int64_t cols;

      ErrorOr<BufferR2<F32>> typed_input =
          Match("input", input,
                m::Buffer<F32>().WithDims(&rows, &cols));
      if (typed_input.has_error()) {
        return std::move(typed_input).error();
      }

      ErrorOr<Result<BufferR2<F32>>> typed_output =
          Match("output", output,
                m::Buffer<F32>().WithDims(rows, cols));
      if (typed_output.has_error()) {
        return std::move(typed_output).error();
      }

      float* input_data = typed_input->typed_data();
      float* output_data = (*typed_output)->typed_data();
      return Error::Success();
    });

이 예시에서 input 일치는 두 측정기준 크기를 캡처하고 output 일치는 동일한 도형을 갖는지 확인합니다. 캡처는 전체 패턴이 성공적으로 일치하는 경우에만 커밋됩니다.

측정기준 패턴은 제약조건이 없거나, 고정되거나, 캡처될 수 있습니다. 고정 값과 캡처 포인터는 측정기준 패턴으로 암시적으로 변환됩니다.

m::Dim()       // Any dimension size.
m::Dim(16)     // A dimension whose size is 16.
m::Dim(&rows)  // Any dimension size, captured in `rows` on success.

m::Buffer<F32>().WithDims(&rows, 16);

WithDims 는 전체 도형을 설명하고 인수 수에서 순위를 설정합니다. WithDim<I> 또는 WithDim(index, ...)은 전체 순위를 고정하지 않고 측정기준을 위치별로 제약합니다. 위치 제약조건에는 참조된 측정기준이 있어야 하지만 다른 모든 측정기준은 제약조건이 없습니다.

ErrorOr<BufferR4<F32>> MatchActivations(AnyBuffer input) {
  return Match(
      "input", input,
      m::Buffer<F32, 4>().WithDim<0>(1).WithDim<3>(128));
}

Error VerifyDimension(AnyBuffer input, size_t index, int64_t size) {
  return Verify("input", input, m::Buffer().WithDim(index, size));
}

MatchActivations 에는 측정기준 0과 3이 각각 1과 128인 순위 4 F32 버퍼가 필요합니다. 측정기준 1과 2는 제한이 없습니다. 명시적인 WithRank가 없으면 WithDim(index, ...)은 index보다 큰 순위를 허용합니다.

측정기준 캡처는 위치 간의 관계를 표현할 수도 있습니다.

auto matrix = m::Buffer<F32>().WithDims(&rows, 128);

int64_t n;
ErrorOr<BufferR2<F32>> square =
    Match("matrix", buffer,
          m::Buffer<F32>().WithDims(&n, &n));

동일한 캡처 포인터가 여러 번 표시되면 모든 상응하는 측정기준의 크기가 동일해야 합니다. 따라서 위의 Match 호출은 정사각형 행렬만 허용합니다. 측정기준 캡처는 전체 패턴이 성공한 후에만 작성되며 실패 시에는 변경되지 않습니다.

전체 버퍼 관계는 모든 측정기준을 캡처할 필요가 없습니다. WithShapeOf 는 다른 dtype을 허용하면서 다른 버퍼의 전체 런타임 도형과 일치합니다. Like에도 동일한 dtype이 필요합니다.

Error VerifySortBuffers(AnyBuffer keys, AnyBuffer values,
                        Result<AnyBuffer> keys_output) {
  Error error =
      Verify("values", values, m::Buffer().WithShapeOf(keys));
  if (error.failure()) {
    return error;
  }

  return Verify("keys output", *keys_output, m::Buffer().Like(keys));
}

두 수정자는 패턴이 빌드될 때 참조 버퍼의 메타데이터를 복사합니다. 패턴은 버퍼에 대한 참조를 유지하지 않습니다. 이는 런타임 제약조건이므로 AnyBuffer를 구체적인 반환 유형으로 세분화하지 않습니다. 이는 주로 Verify 또는 입력 버퍼에 이미 구체적인 유형이 있는 경우 Match와 함께 유용합니다.

버퍼가 기존 유형을 유지해야 하거나 패턴이 둘 이상의 데이터 유형 또는 순위를 허용하는 경우 Verify를 사용합니다. 예를 들어 색인 버퍼는 S32 또는 S64를 허용하고 확인 후 데이터 유형에 따라 디스패치할 수 있습니다.

namespace m = ::xla::ffi::match;

Error VerifyIndices(AnyBuffer indices) {
  auto pattern = m::Buffer().WithDType<S32, S64>().WithRank<1, 2>();

  Error error = Verify("indices", indices, pattern);
  if (error.failure()) {
    return error;
  }

  // Dispatch an implementation based on `indices.element_type()`.
  return Error::Success();
}

이 패턴은 여러 구체적인 버퍼 유형을 허용하므로 Match로 AnyBuffer를 세분화할 수 없습니다. 고유한 반환 유형이 없습니다. 반환 유형이 이미 알려진 이미 유형이 지정된 버퍼를 사용하여 Match에 전달할 수 있습니다. 이미 유형이 지정된 버퍼도 확인할 수 있습니다. 이는 일반적으로 도형 관계를 확인하는 데 유용합니다.

Error VerifySameShape(BufferR2<F32> lhs, BufferR2<F32> rhs) {
  int64_t rows;
  int64_t cols;

  Error error = Verify(
      "lhs", lhs,
      m::Buffer().WithDims(&rows, &cols));
  if (error.failure()) {
    return error;
  }

  return Verify("rhs", rhs,
                m::Buffer().WithDims(rows, cols));
}

Match 및 Verify에 첫 번째 인수로 전달된 이름은 필수이며 실패한 제약조건에 대한 설명과 함께 오류에 포함됩니다.

xla/ffi/api/ffi.h의 외부 FFI API는 위의 예와 같이 Verify에서 Error를 반환하고 Match에서 ErrorOr<T>를 반환합니다. 내부 FFI API는 xla/ffi/ffi.h를 사용하여 동일한 일치 인터페이스를 제공하며 absl::Status 및 absl::StatusOr<T>를 사용합니다.

ABSL_ASSIGN_OR_RETURN(BufferR2<F32> input,
                 Match("input", buffer, m::Buffer<F32, 2>()));
ABSL_RETURN_IF_ERROR(Verify(
    "input", input,
    m::Buffer().WithDims(expected_rows, m::Dim())));

가변 인수 및 결과

커스텀 호출의 여러 인스턴스에서 인수 및 결과 수가 다를 수 있는 경우 RemainingArgs 및 RemainingRets를 사용하여 런타임에 디코딩할 수 있습니다.

auto handler = Ffi::Bind().RemainingArgs().RemainingRets().To(
    [](RemainingArgs args, RemainingRets results) -> Error {
      ErrorOr<AnyBuffer> arg = args.get<AnyBuffer>(0);
      ErrorOr<Result<AnyBuffer>> res = results.get<AnyBuffer>(0);

      if (!arg.has_value()) {
        return Error(ErrorCode::kInternal, arg.error());
      }

      if (!res.has_value()) {
        return Error(ErrorCode::kInternal, res.error());
      }

      return Error::Success();
    });

가변 인수 및 결과는 일반 인수 및 결과 뒤에 선언할 수 있지만 가변 인수 뒤에 일반 인수 및 결과를 결합하는 것은 허용되지 않습니다.

auto handler =
    Ffi::Bind()
        .Arg<AnyBuffer>()
        .RemainingArgs()
        .Ret<AnyBuffer>()
        .RemainingRets()
        .To([](AnyBuffer arg, RemainingArgs args, AnyBuffer ret,
               RemainingRets results) -> Error { return Error::Success(); });

속성

XLA FFI는 custom_call backend_config로 전달된 mlir::DictionaryAttr를 FFI 핸들러 인수로 자동 디코딩하는 것을 지원합니다.

%0 = "stablehlo.custom_call"(%arg0) {
  call_target_name = "foo",
  backend_config= {
    i32 = 42 : i32,
    str = "string"
  },
  api_version = 4 : i32
} : (tensor<f32>) -> tensor<f32>

이 예시에서 커스텀 호출에는 단일 버퍼 인수와 두 개의 속성이 있으며, XLA FFI는 이를 자동으로 디코딩하여 사용자 정의 호출 가능 항목에 전달할 수 있습니다.

auto handler = Ffi::Bind()
  .Arg<BufferR0<F32>>()
  .Attr<int32_t>("i32")
  .Attr<std::string_view>("str")
  .To([](BufferR0<F32> buffer, int32_t i32, std::string_view str) {
    return Error::Success();
  });

사용자 정의 열거형 속성

XLA FFI는 정수 MLIR 속성을 사용자 정의 열거형으로 자동 디코딩할 수 있습니다. 열거형 클래스는 동일한 기본 정수 유형을 가져야 하며 디코딩은 XLA FFI에 명시적으로 등록되어야 합니다.

%0 = "stablehlo.custom_call"(%arg0) {
  call_target_name = "foo",
  backend_config= {
    command = 0 : i32
  },
  api_version = 4 : i32
} : (tensor<f32>) -> tensor<f32>

enum class Command : int32_t {
  kAdd = 0,
  kMul = 1,
};

XLA_FFI_REGISTER_ENUM_ATTR_DECODING(Command);

auto handler = Ffi::Bind().Attr<Command>("command").To(
    [](Command command) -> Error { return Error::Success(); });

모든 커스텀 호출 속성 결합

모든 커스텀 호출 속성에 사전으로 액세스하고 런타임에 필요한 속성만 지연 디코딩할 수 있습니다.

auto handler = Ffi::Bind().Attrs().To([](Dictionary attrs) -> Error {
  ErrorOr<int32_t> i32 = attrs.get<int32_t>("i32");
  return Error::Success();
});

사용자 정의 구조체 속성

XLA FFI는 사전 속성을 사용자 정의 구조체로 디코딩할 수 있습니다.

%0 = "stablehlo.custom_call"(%arg0) {
  call_target_name = "foo",
  backend_config= {
    range = { lo = 0 : i64, hi = 42 : i64 }
  },
  api_version = 4 : i32
} : (tensor<f32>) -> tensor<f32>

위의 예에서 range는 mlir::DictionaryAttr 속성이며, 이름으로 사전 필드에 액세스하는 대신 C++ 구조체로 자동 디코딩할 수 있습니다. 디코딩은 XLA_FFI_REGISTER_STRUCT_ATTR_DECODING 매크로에 명시적으로 등록되어야 합니다. 백그라운드에서 ::xla::ffi 네임스페이스에 템플릿 특수화를 정의하므로 매크로를 전역 네임스페이스에 추가해야 합니다.

struct Range {
  int64_t lo;
  int64_t hi;
};

XLA_FFI_REGISTER_STRUCT_ATTR_DECODING(Range, StructMember<int64_t>("lo"),
                                             StructMember<int64_t>("hi"));

auto handler = Ffi::Bind().Attr<Range>("range").To([](Range range) -> Error{
  return Error::Success();
});

커스텀 속성은 다른 속성과 마찬가지로 사전에서 로드할 수 있습니다. 아래 예에서 모든 커스텀 호출 속성은 Dictionary로 디코딩되며 range는 이름으로 액세스할 수 있습니다.

auto handler = Ffi::Bind().Attrs().To([](Dictionary attrs) -> Error {
  ErrorOr<Range> range = attrs.get<Range>("range");
  return Error::Success();
});

CPU에서 커스텀 호출 만들기

XLA의 클라이언트 API를 통해 커스텀 호출을 나타내는 HLO 명령어를 만들 수 있습니다. 예를 들어 다음 코드는 커스텀 호출을 사용하여 CPU에서 A[i] = B[i % 128]+ C[i]를 계산합니다. (물론 일반 HLO로 이 작업을 할 수 있으며 해야 합니다. )

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

void do_it() {
  xla::XlaBuilder b("do_it");
  xla::XlaOp param0 =
      xla::Parameter(&b, 0, xla::ShapeUtil::MakeShape(xla::F32, {128}), "p0");
  xla::XlaOp param1 =
      xla::Parameter(&b, 1, xla::ShapeUtil::MakeShape(xla::F32, {2048}), "p1");
  xla::XlaOp custom_call =
      xla::CustomCall(&b, "do_custom_call", /*operands=*/{param0, param1},
        /*shape=*/xla::ShapeUtil::MakeShape(xla::F32, {2048}),
        /*opaque=*/"", /*has_side_effect=*/false,
        /*output_operand_aliasing=*/{}, /*literal=*/nullptr,
        /*schedule=*/CustomCallSchedule::SCHEDULE_NONE,
        /*api_version=*/CustomCallApiVersion::API_VERSION_TYPED_FFI);
}

// Constrain custom call arguments to 1-dimensional buffers of F32 data type.
using BufferF32 = xla::ffi::BufferR1<xla::ffi::DataType::F32>;

// Implement a custom call as a C++ function. Note that we can use `Buffer` type
// defined by XLA FFI that gives us access to buffer data type and shape.
xla::ffi::Error do_custom_call(BufferF32 in0, BufferF32 in1,
                               xla::ffi::Result<BufferF32> out) {
  size_t d0 = in0.dimensions[0];
  size_t d1 = in1.dimensions[0];

  // Check that dimensions are compatible.
  assert(out->dimensions[0] == d1 && "unexpected dimensions");

  for (size_t i = 0; i < d1; ++i) {
    out->data[i] = in0.data[i % d0] + in1.data[i];
  }
}

// Explicitly define an XLA FFI handler signature and bind it to the
// `do_custom_call` implementation. XLA FFI handler can automatically infer
// type signature from the custom call function, but it relies on magical
// template metaprogramming an explicit binding provides and extra level of
// type checking and clearly states custom call author intentions.
XLA_FFI_DEFINE_HANDLER(handler, do_custom_call,
                       ffi::Ffi::Bind()
                           .Arg<Buffer>()
                           .Arg<Buffer>()
                           .Ret<Buffer>());

// Registers `handler` with and XLA FFI on a "Host" platform.
XLA_FFI_REGISTER_HANDLER(xla::ffi::GetXlaFfiApi(), "do_custom_call",
                         "Host", handler);

GPU에서 커스텀 호출 만들기

XLA FFI를 사용한 GPU 커스텀 호출 등록은 거의 동일하며, 유일한 차이점은 GPU의 경우 기기에서 커널을 실행할 수 있도록 기본 플랫폼 스트림(CUDA 또는 ROCM 스트림)을 요청해야 한다는 것입니다. 다음은 위의 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 = blockIdx.x * blockDim.x + threadIdx.x;
  out[idx] = in0[idx % 128] + in1[idx];
}

void do_custom_call(CUstream stream, BufferF32 in0, BufferF32 in1,
                    xla::ffi::Result<BufferF32> out) {
  size_t d0 = in0.dimensions[0];
  size_t d1 = in1.dimensions[0];
  size_t d2 = out->dimensions[0];

  assert(d0 == 128 && d1 == 2048 && d2 == 2048 && "unexpected dimensions");

  const int64_t block_dim = 64;
  const int64_t grid_dim = 2048 / block_dim;
  custom_call_kernel<<<grid_dim, block_dim, 0, stream>>>(
    in0.data, in1.data, out->data);
}

XLA_FFI_DEFINE_HANDLER(handler, do_custom_call,
                       ffi::Ffi::Bind()
                           .Ctx<xla::ffi::PlatformStream<CUstream>>()
                           .Arg<BufferF32>()
                           .Arg<BufferF32>()
                           .Ret<BufferF32>());

XLA_FFI_REGISTER_HANDLER(xla::ffi::GetXlaFfiApi(), "do_custom_call",
                         "CUDA", handler);

먼저 GPU 커스텀 호출 함수는 여전히 CPU에서 실행되는 함수 라는 점에 유의하세요. do_custom_call CPU 함수는 GPU에서 작업을 대기열에 추가하는 역할을 합니다. 여기서는 CUDA 커널을 실행하지만 cuBLAS 호출과 같은 다른 작업을 실행할 수도 있습니다.

인수와 결과도 호스트에 있으며 데이터 멤버에는 기기 (즉, GPU) 메모리에 대한 포인터가 포함되어 있습니다. 커스텀 호출 핸들러에 전달된 버퍼는 기본 기기 버퍼의 도형을 가지므로 커스텀 호출은 이를 통해 커널 실행 매개변수를 계산할 수 있습니다.

커스텀 호출에 튜플 전달

다음 커스텀 호출을 살펴보세요.

using xla::ShapeUtil;
using xla::F32;
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 모두에서 튜플은 메모리에 포인터 배열로 표시됩니다. XLA가 튜플 인수 또는 결과로 커스텀 호출을 호출하면 이를 평면화하고 일반 버퍼 인수 또는 결과로 전달합니다.

임시 버퍼로 튜플 출력

커스텀 호출에 대한 튜플 입력은 편리하지만 엄격히 말하면 필요하지 않습니다. 커스텀 호출에 대한 튜플 입력을 지원하지 않는 경우 항상 get-tuple-element를 사용하여 튜플을 압축 해제한 후 커스텀 호출에 전달할 수 있습니다.

반면에 튜플 출력 을 사용하면 다른 방법으로는 할 수 없는 작업을 할 수 있습니다.

튜플 출력이 있는 명백한 이유는 튜플 출력이 커스텀 호출 (또는 다른 XLA 작업)이 여러 독립 배열을 반환하는 방법이기 때문입니다.

하지만 튜플 출력은 커스텀 호출 임시 메모리를 제공하는 방법이기도 합니다. 예, 출력 은 임시 버퍼를 나타낼 수 있습니다. 출력 버퍼에는 작업이 쓸 수 있고 쓰여진 후에 읽을 수 있는 속성이 있습니다. 이것이 바로 임시 버퍼에서 원하는 것입니다.

위의 예에서 F32[1024]를 임시 버퍼로 사용한다고 가정해 보겠습니다. 그런 다음 위의 HLO를 작성하고 커스텀 호출 출력의 튜플 색인 1을 읽지 않습니다.