XLA कस्टम कॉल

इस दस्तावेज़ में, XLA FFI लाइब्रेरी का इस्तेमाल करके, XLA कस्टम कॉल लिखने और इस्तेमाल करने का तरीका बताया गया है. कस्टम कॉल, XLA कंपाइलर को कंपाइल करने के समय, HLO मॉड्यूल में बाहरी "ऑपरेशन" के बारे में बताने का एक तरीका है. वहीं, XLA FFI, XLA के साथ ऐसे ऑपरेशन को रजिस्टर करने का एक तरीका है. यह काम रन टाइम पर होता है. FFI का मतलब "फ़ॉरेन फ़ंक्शन इंटरफ़ेस" है. यह C एपीआई का एक सेट है. यह XLA के लिए एक बाइनरी इंटरफ़ेस (ABI) तय करता है, ताकि अन्य प्रोग्रामिंग भाषाओं में लिखे गए बाहरी कोड को कॉल किया जा सके. XLA, C++ में लिखे गए XLA FFI के लिए सिर्फ़ हेडर बाइंडिंग उपलब्ध कराता है. इससे, C API की सभी जानकारी, उपयोगकर्ता से छिपी रहती है.

JAX + XLA कस्टम कॉल

कस्टम कॉल और XLA FFI को JAX के साथ इंटिग्रेट करने के एंड-टू-एंड उदाहरणों के लिए, JAX का दस्तावेज़ देखें.

XLA FFI बाइंडिंग

XLA FFI बाइंडिंग, कंपाइल-टाइम में कस्टम कॉल सिग्नेचर की खास जानकारी होती है: कस्टम कॉल के तर्क, एट्रिब्यूट, और उनके टाइप, और एक्ज़ीक्यूशन कॉन्टेक्स्ट के ज़रिए पास किए गए अतिरिक्त पैरामीटर (जैसे, GPU बैकएंड के लिए GPU स्ट्रीम). XLA FFI बाइंडिंग को, C++ के किसी भी कॉल किए जा सकने वाले फ़ंक्शन (फ़ंक्शन पॉइंटर, लैम्डा वगैरह) से बाइंड किया जा सकता है. इसके लिए, operator() का सिग्नेचर मेल खाना चाहिए. कंस्ट्रक्टेड हैंडलर, XLA FFI कॉल फ़्रेम को डिकोड करता है. इसे स्टेबल सी एपीआई से तय किया जाता है. साथ ही, यह सभी पैरामीटर की टाइप जांच करता है और डिकोड किए गए नतीजों को उपयोगकर्ता के तय किए गए कॉलबैक पर फ़ॉरवर्ड करता है.

XLA FFI बाइंडिंग, टेंप्लेट मेटाप्रोग्रामिंग पर काफ़ी हद तक निर्भर करती है. इससे कंस्ट्रक्ट किए गए हैंडलर को सबसे असरदार मशीन कोड में कंपाइल किया जा सकता है. हर कस्टम कॉल पैरामीटर के लिए, रन टाइम ओवरहेड कुछ नैनोसेकंड के क्रम में होते हैं.

XLA FFI के कस्टम पॉइंट, टेंप्लेट स्पेशलाइज़ेशन के तौर पर लागू किए जाते हैं.साथ ही, उपयोगकर्ता यह तय कर सकते हैं कि उनके कस्टम टाइप को कैसे डिकोड किया जाए. इसका मतलब है कि उपयोगकर्ता के तय किए गए enum class टाइप के लिए, कस्टम डिकोडिंग तय की जा सकती है.

कस्टम कॉल से गड़बड़ियां वापस लाना

कस्टम कॉल लागू करने पर, XLA रनटाइम को xla::ffi::Error वैल्यू दिखनी चाहिए, ताकि उसे सफलता या गड़बड़ी के बारे में पता चल सके. यह absl::Status के जैसा ही है और इसमें गड़बड़ी के कोड का एक ही सेट होता है. हम absl::Status का इस्तेमाल नहीं करते, क्योंकि इसमें स्टेबल एबीआई नहीं है. साथ ही, इसे डाइनैमिक तौर पर लोड की गई कस्टम कॉल लाइब्रेरी और 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 के लिए, रैंक-4 F32 बफ़र की ज़रूरत होती है. इसके डाइमेंशन 0 और 3, क्रमशः 1 और 128 के बराबर होने चाहिए. डाइमेंशन 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 एपीआई, Verify से Error और Match से ErrorOr<T> दिखाता है. जैसा कि ऊपर दिए गए उदाहरणों में दिखाया गया है. xla/ffi/ffi.h में मौजूद इंटरनल FFI API, 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, mlir::DictionaryAttr को custom_call backend_config के तौर पर पास किए जाने पर, उसे 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();
  });

उपयोगकर्ता के तय किए गए enum एट्रिब्यूट

XLA FFI, इंटिग्रल MLIR एट्रिब्यूट को उपयोगकर्ता की तय की गई enum में अपने-आप डिकोड कर सकता है. एनम क्लास का इंटिग्रल टाइप एक ही होना चाहिए. साथ ही, डिकोडिंग को 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();
});

कस्टम एट्रिब्यूट को किसी शब्दकोश से लोड किया जा सकता है. ठीक वैसे ही जैसे किसी अन्य एट्रिब्यूट को लोड किया जाता है. नीचे दिए गए उदाहरण में, सभी कस्टम कॉल एट्रिब्यूट को a Dictionary के तौर पर डिकोड किया गया है. साथ ही, a range को नाम के हिसाब से ऐक्सेस किया जा सकता है.

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

सीपीयू पर कस्टम कॉल बनाना

XLA के क्लाइंट एपीआई के ज़रिए, कस्टम कॉल को दिखाने वाला एचएलओ निर्देश बनाया जा सकता है. उदाहरण के लिए, इस कोड में सीपीयू पर A[i] = B[i % 128]+ C[i] को कंप्यूट करने के लिए, कस्टम कॉल का इस्तेमाल किया गया है. (बिलकुल, ऐसा किया जा सकता है और करना भी चाहिए! – do this with regular 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);

जीपीयू पर कस्टम कॉल बनाना

XLA FFI के साथ जीपीयू कस्टम कॉल रजिस्ट्रेशन लगभग एक जैसा होता है. सिर्फ़ यह अंतर होता है कि जीपीयू के लिए, आपको डिवाइस पर कर्नल लॉन्च करने के लिए, प्लैटफ़ॉर्म स्ट्रीम (CUDA या ROCM स्ट्रीम) के बारे में पूछना होता है. यहां 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 = 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 कस्टम कॉल फ़ंक्शन अब भी सीपीयू पर एक्ज़ीक्यूट होने वाला फ़ंक्शन है. do_custom_call सीपीयू फ़ंक्शन, जीपीयू पर काम को कतार में लगाने के लिए ज़िम्मेदार होता है. यहां यह CUDA कर्नल लॉन्च करता है. हालांकि, यह कुछ और भी कर सकता है, जैसे कि cuBLAS को कॉल करना.

आर्ग्युमेंट और नतीजे भी होस्ट पर मौजूद होते हैं. साथ ही, डेटा मेंबर में डिवाइस (यानी कि जीपीयू) की मेमोरी का पॉइंटर होता है. कस्टम कॉल हैंडलर को पास किए गए बफ़र, डिवाइस के बफ़र के आकार के होते हैं. इसलिए, कस्टम कॉल उनसे कर्नल लॉन्च पैरामीटर का हिसाब लगा सकता है.

कस्टम कॉल में टपल पास करना

यहां कस्टम कॉल का एक उदाहरण दिया गया है.

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, ...);

सीपीयू और जीपीयू, दोनों पर टपल को मेमोरी में पॉइंटर की एक कैटगरी के तौर पर दिखाया जाता है. जब XLA, टपल आर्ग्युमेंट या नतीजों के साथ कस्टम कॉल करता है, तो यह उन्हें फ़्लैट कर देता है और सामान्य बफ़र आर्ग्युमेंट या नतीजों के तौर पर पास करता है.

टपल को अस्थायी बफ़र के तौर पर आउटपुट करता है

कस्टम कॉल के लिए टपल इनपुट का इस्तेमाल करना आसान होता है, लेकिन यह ज़रूरी नहीं है. अगर हम कस्टम कॉल के लिए टपल इनपुट का इस्तेमाल नहीं करते हैं, तो कस्टम कॉल को पास करने से पहले, get-tuple-element का इस्तेमाल करके टपल को हमेशा अनपैक किया जा सकता है.

दूसरी ओर, टपल आउटपुट से ऐसे काम किए जा सकते हैं जो किसी और तरीके से नहीं किए जा सकते.

टपल आउटपुट का इस्तेमाल करने की मुख्य वजह यह है कि कस्टम कॉल (या कोई अन्य XLA ऑप) कई इंडिपेंडेंट ऐरे को टपल आउटपुट के तौर पर दिखाता है.

हालांकि, यह कम ही लोगों को पता होगा कि टपल आउटपुट, कस्टम कॉल टेंप मेमोरी देने का भी एक तरीका है. हां, आउटपुट, टेंप बफ़र को दिखा सकता है. मान लें कि किसी आउटपुट बफ़र में यह प्रॉपर्टी होती है कि ऑप इसमें लिख सकता है. साथ ही, इसमें लिखने के बाद इसे पढ़ा जा सकता है. टेम्प बफ़र से आपको यही चाहिए.

ऊपर दिए गए उदाहरण में, मान लें कि हमें F32[1024] को अस्थायी बफ़र के तौर पर इस्तेमाल करना है. इसके बाद, हम ऊपर की तरह ही एचएलओ लिखेंगे. साथ ही, हम कस्टम कॉल के आउटपुट के टपल इंडेक्स 1 को कभी नहीं पढ़ेंगे.