Dokumen ini menjelaskan cara menulis dan menggunakan panggilan kustom XLA menggunakan library XLA FFI. Panggilan kustom adalah mekanisme untuk mendeskripsikan "operasi" eksternal dalam modul HLO ke compiler XLA (pada waktu kompilasi), dan XLA FFI adalah mekanisme untuk mendaftarkan implementasi operasi tersebut dengan XLA (pada waktu runtime). FFI adalah singkatan dari "foreign function interface" dan merupakan kumpulan C API yang menentukan antarmuka biner (ABI) untuk XLA memanggil kode eksternal yang ditulis dalam bahasa pemrograman lain. XLA menyediakan binding khusus header untuk XLA FFI yang ditulis dalam C++, yang menyembunyikan semua detail tingkat rendah dari C API yang mendasarinya dari pengguna akhir.
Panggilan Kustom JAX + XLA
Lihat dokumentasi JAX untuk contoh menyeluruh tentang mengintegrasikan panggilan kustom dan XLA FFI dengan JAX.
Binding XLA FFI
Binding XLA FFI adalah spesifikasi waktu kompilasi dari tanda tangan panggilan kustom: argumen panggilan kustom, atribut dan jenisnya, serta parameter tambahan yang diteruskan melalui konteks eksekusi (yaitu, aliran GPU untuk backend GPU). Binding XLA FFI dapat terikat ke C++ callable (pointer fungsi, lambda, dll.) dengan tanda tangan operator() yang kompatibel. Handler yang dibuat mendekode frame panggilan XLA FFI (ditentukan oleh C API yang stabil), memeriksa jenis semua parameter, dan meneruskan hasil yang didekode ke callback yang ditentukan pengguna.
Binding XLA FFI sangat bergantung pada metaprogramming template agar dapat mengompilasi handler yang dibuat ke kode mesin yang paling efisien. Overhead waktu proses berada dalam urutan beberapa nanodetik untuk setiap parameter panggilan kustom.
Titik penyesuaian XLA FFI diimplementasikan sebagai spesialisasi template, dan pengguna dapat menentukan cara mendekode jenis kustom mereka, yaitu, memungkinkan untuk menentukan decoding kustom untuk jenis enum class yang ditentukan pengguna.
Menampilkan Error Dari Panggilan Kustom
Implementasi panggilan kustom harus menampilkan nilai xla::ffi::Error untuk menandakan keberhasilan atau error ke runtime XLA. Nilai ini mirip dengan absl::Status, dan memiliki kumpulan kode error yang sama. Kami tidak menggunakan absl::Status karena tidak memiliki ABI yang stabil dan tidak aman untuk meneruskannya antara library panggilan kustom yang dimuat secara dinamis, dan XLA itu sendiri.
// 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(); });
Argumen dan Hasil Buffer
XLA menggunakan gaya penerusan tujuan untuk hasil: panggilan kustom (atau operasi XLA lainnya) tidak mengalokasikan memori untuk hasil, dan menulis ke tujuan yang diteruskan oleh runtime XLA. XLA menggunakan penetapan buffer statis, dan mengalokasikan buffer untuk semua nilai berdasarkan rentang aktifnya pada waktu kompilasi.
Hasil yang diteruskan ke handler FFI digabungkan ke dalam Result<T> template, yang
memiliki semantik seperti pointer: operator-> memberikan akses ke parameter yang mendasarinya.
Argumen dan hasil AnyBuffer memberikan akses ke parameter buffer panggilan kustom dari jenis data apa pun. Hal ini berguna saat panggilan kustom memiliki implementasi generik yang berfungsi untuk beberapa jenis data, dan implementasi panggilan kustom menjalankan pengiriman waktu proses berdasarkan jenis data. AnyBuffer memberikan akses ke jenis data buffer, dimensi, dan pointer ke buffer itu sendiri.
%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();
});
Argumen dan Hasil Buffer Terbatas
Buffer memungkinkan untuk menambahkan batasan pada jenis data buffer dan jumlah dimensi, dan batasan tersebut akan otomatis diperiksa oleh handler dan menampilkan error ke runtime XLA, jika argumen waktu proses tidak cocok dengan tanda tangan handler FFI.
// 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();
});
Mencocokkan dan Memverifikasi Buffer
Pola buffer dapat mengekspresikan batasan yang lebih spesifik daripada jenis argumen dan hasil dalam binding FFI. Pola ini berguna untuk menyempurnakan AnyBuffer ke jenis buffer konkret, memeriksa ukuran dimensi, dan memeriksa hubungan antara beberapa bentuk buffer.
Pola bersifat tidak dapat diubah dan dapat menentukan dtype dan peringkat baik secara langsung maupun dengan pengubah yang sesuai:
namespace m = ::xla::ffi::match;
m::Buffer<F32, 2>();
m::Buffer().WithDType<F32>().WithRank<2>();
Fungsi Match memeriksa pola dan menampilkan jenis buffer yang dibatasi secara maksimal. Untuk menyempurnakan AnyBuffer, pola harus menentukan tepat satu jenis data dan satu peringkat. Mencocokkan buffer yang sudah diketik akan memverifikasi batasan tambahan dan mempertahankan jenisnya.
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();
});
Dalam contoh ini, pencocokan input mengambil dua ukuran dimensinya, dan pencocokan output memverifikasi bahwa output memiliki bentuk yang sama. Pengambilan hanya dilakukan saat pola lengkap berhasil dicocokkan.
Pola dimensi dapat tidak dibatasi, tetap, atau diambil. Nilai tetap dan pointer pengambilan secara implisit dikonversi ke pola dimensi:
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 menjelaskan bentuk lengkap dan menetapkan peringkatnya dari jumlah argumen. WithDim<I> atau WithDim(index, ...) membatasi dimensi
secara posisional tanpa memperbaiki peringkat lengkap. Batasan posisional mengharuskan dimensi yang direferensikan ada, tetapi membiarkan setiap dimensi lainnya tidak dibatasi:
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 memerlukan buffer F32 peringkat 4 dengan dimensi 0 dan 3 sama dengan 1 dan 128, masing-masing; dimensi 1 dan 2 tidak dibatasi. Tanpa
eksplisit WithRank, WithDim(index, ...) menerima peringkat apa pun yang lebih besar dari
index.
Pengambilan dimensi juga dapat mengekspresikan hubungan antar-posisi:
auto matrix = m::Buffer<F32>().WithDims(&rows, 128);
int64_t n;
ErrorOr<BufferR2<F32>> square =
Match("matrix", buffer,
m::Buffer<F32>().WithDims(&n, &n));
Jika pointer pengambilan yang sama muncul beberapa kali, semua dimensi yang sesuai harus memiliki ukuran yang sama. Oleh karena itu, panggilan Match di atas hanya menerima matriks persegi. Pengambilan dimensi hanya ditulis setelah pola lengkap berhasil dan tidak berubah jika gagal.
Hubungan seluruh buffer tidak memerlukan pengambilan setiap dimensi.
WithShapeOf cocok dengan bentuk runtime lengkap dari buffer lain sekaligus memungkinkan dtype yang berbeda. Like juga memerlukan dtype yang sama:
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));
}
Kedua pengubah menyalin metadata buffer referensi saat pola dibuat; pola tidak mempertahankan referensi ke buffer. Karena ini adalah batasan runtime, batasan ini tidak menyempurnakan AnyBuffer ke jenis nilai yang ditampilkan konkret. Batasan ini terutama berguna dengan Verify, atau dengan Match saat buffer input sudah memiliki jenis konkret.
Gunakan Verify saat buffer harus mempertahankan jenis yang ada, atau saat pola menerima lebih dari satu jenis data atau peringkat. Misalnya, buffer indeks dapat menerima S32 atau S64 dan mengirim berdasarkan jenis datanya setelah verifikasi:
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();
}
Karena pola ini memungkinkan beberapa jenis buffer konkret, pola ini tidak dapat menyempurnakan AnyBuffer dengan Match: tidak ada jenis nilai yang ditampilkan yang unik. Pola ini masih dapat diteruskan ke Match dengan buffer yang sudah diketik, yang jenis nilai yang ditampilkan sudah diketahui. Buffer yang sudah diketik juga dapat diverifikasi; hal ini biasanya berguna untuk memeriksa hubungan bentuk:
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));
}
Nama yang diteruskan sebagai argumen pertama ke Match dan Verify diperlukan dan disertakan dalam error, bersama dengan deskripsi batasan yang gagal.
FFI API eksternal di xla/ffi/api/ffi.h menampilkan Error dari Verify dan
ErrorOr<T> dari Match, seperti dalam contoh di atas. FFI API internal di
xla/ffi/ffi.h menyediakan antarmuka pencocokan yang sama menggunakan absl::Status dan
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())));
Argumen dan Hasil Variadik
Jika jumlah argumen dan hasil dapat berbeda dalam instance panggilan kustom yang berbeda, argumen dan hasil tersebut dapat didekode pada waktu proses menggunakan RemainingArgs dan 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();
});
Argumen dan hasil variadik dapat dideklarasikan setelah argumen dan hasil reguler, tetapi mengikat argumen dan hasil reguler setelah argumen dan hasil variadik adalah ilegal.
auto handler =
Ffi::Bind()
.Arg<AnyBuffer>()
.RemainingArgs()
.Ret<AnyBuffer>()
.RemainingRets()
.To([](AnyBuffer arg, RemainingArgs args, AnyBuffer ret,
RemainingRets results) -> Error { return Error::Success(); });
Atribut
XLA FFI mendukung decoding otomatis mlir::DictionaryAttr yang diteruskan sebagai custom_call backend_config ke dalam argumen handler 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>
Dalam contoh ini, panggilan kustom memiliki satu argumen buffer dan dua atribut, dan XLA FFI dapat otomatis mendekodenya dan meneruskan ke callable yang ditentukan pengguna.
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();
});
Atribut Enum yang Ditentukan Pengguna
XLA FFI dapat otomatis mendekode atribut MLIR integral ke dalam enum yang ditentukan pengguna. Class enum harus memiliki jenis integral yang mendasarinya, dan decoding harus didaftarkan secara eksplisit dengan 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(); });
Mengikat Semua Atribut Panggilan Kustom
Anda dapat mengakses semua atribut panggilan kustom sebagai kamus dan mendekode secara lambat hanya atribut yang diperlukan pada waktu proses.
auto handler = Ffi::Bind().Attrs().To([](Dictionary attrs) -> Error {
ErrorOr<int32_t> i32 = attrs.get<int32_t>("i32");
return Error::Success();
});
Atribut Struct yang Ditentukan Pengguna
XLA FFI dapat mendekode atribut kamus ke dalam struct yang ditentukan pengguna.
%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>
Dalam contoh di atas, range adalah atribut mlir::DictionaryAttr, dan bukan mengakses kolom kamus berdasarkan nama, kolom tersebut dapat otomatis didekode sebagai struct C++. Decoding harus didaftarkan secara eksplisit dengan makro XLA_FFI_REGISTER_STRUCT_ATTR_DECODING (di balik layar, decoding menentukan spesialisasi template di namespace ::xla::ffi, sehingga makro harus ditambahkan ke namespace global).
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();
});
Atribut kustom dapat dimuat dari kamus, seperti atribut lainnya. Dalam contoh di bawah, semua atribut panggilan kustom didekode sebagai Dictionary, dan range dapat diakses berdasarkan nama.
auto handler = Ffi::Bind().Attrs().To([](Dictionary attrs) -> Error {
ErrorOr<Range> range = attrs.get<Range>("range");
return Error::Success();
});
Membuat panggilan kustom di CPU
Anda dapat membuat instruksi HLO yang mewakili panggilan kustom melalui client API XLA. Misalnya, kode berikut menggunakan panggilan kustom untuk menghitung A[i] = B[i %
128]+ C[i] di CPU. (Tentu saja Anda dapat – dan harus! – melakukannya dengan HLO reguler.)
#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);
Membuat panggilan kustom di GPU
Pendaftaran panggilan kustom GPU dengan XLA FFI hampir identik, satu-satunya perbedaan adalah bahwa untuk GPU, Anda harus meminta aliran platform yang mendasarinya (aliran CUDA atau ROCM) agar dapat meluncurkan kernel di perangkat. Berikut adalah contoh CUDA yang melakukan komputasi yang sama (A[i] = B[i % 128] + C[i]) seperti kode CPU di atas.
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);
Perhatikan terlebih dahulu bahwa fungsi panggilan kustom GPU masih merupakan fungsi yang dieksekusi di CPU. Fungsi CPU do_custom_call bertanggung jawab untuk mengantrekan pekerjaan di GPU. Di sini, fungsi ini meluncurkan kernel CUDA, tetapi juga dapat melakukan hal lain, seperti memanggil cuBLAS.
Argumen dan hasil juga berada di host, dan anggota data berisi pointer ke memori perangkat (yaitu GPU). Buffer yang diteruskan ke handler panggilan kustom memiliki bentuk buffer perangkat yang mendasarinya, sehingga panggilan kustom dapat menghitung parameter peluncuran kernel dari buffer tersebut.
Meneruskan tuple ke panggilan kustom
Pertimbangkan panggilan kustom berikut.
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, ...);
Di CPU dan GPU, tuple direpresentasikan dalam memori sebagai array pointer. Saat XLA memanggil panggilan kustom dengan argumen atau hasil tuple, XLA akan meratakan argumen atau hasil tersebut dan meneruskannya sebagai argumen atau hasil buffer reguler.
Output tuple sebagai buffer sementara
Input tuple ke panggilan kustom adalah kemudahan, tetapi tidak benar-benar diperlukan. Jika kami tidak mendukung input tuple ke panggilan kustom, Anda selalu dapat mengekstrak tuple menggunakan get-tuple-element sebelum meneruskannya ke panggilan kustom.
Di sisi lain, output tuple memungkinkan Anda melakukan hal-hal yang tidak dapat Anda lakukan.
Alasan yang jelas untuk memiliki output tuple adalah bahwa output tuple adalah cara panggilan kustom (atau operasi XLA lainnya) menampilkan beberapa array independen.
Namun, yang kurang jelas, output tuple juga merupakan cara untuk memberikan memori sementara panggilan kustom Anda. Ya, output dapat mewakili buffer sementara. Pertimbangkan, buffer output memiliki properti yang dapat ditulis oleh operasi, dan dapat dibaca setelah ditulis. Itulah yang Anda inginkan dari buffer sementara.
Dalam contoh di atas, misalkan kita ingin menggunakan F32[1024] sebagai buffer sementara.
Kemudian, kita akan menulis HLO seperti di atas, dan kita tidak akan pernah membaca indeks tuple 1 dari output panggilan kustom.