本文說明如何使用 XLA FFI 程式庫，編寫及使用 XLA 自訂呼叫。自訂呼叫是一種機制，可在 HLO 模組中向 XLA 編譯器 (在編譯時間) 說明外部「作業」，而 XLA FFI 則是一種機制，可在 XLA 中註冊這類作業的實作 (在執行階段)。FFI 是「外部函式介面」的縮寫，是一組 C API，可為 XLA 定義二進位介面 (ABI)，以便呼叫以其他程式設計語言編寫的外部程式碼。XLA 提供以 C++ 編寫的 XLA FFI 專屬標頭繫結，可向使用者隱藏基礎 C API 的所有低階詳細資料。
JAX + XLA 自訂呼叫
如需整合自訂呼叫和 XLA FFI 與 JAX 的端對端範例，請參閱 JAX 說明文件。
XLA FFI 繫結
XLA FFI 繫結是自訂呼叫簽章的編譯時間規格：自訂呼叫引數、屬性和其類型，以及透過執行環境傳遞的其他參數 (即 GPU 後端的 GPU 串流)。XLA FFI 繫結可繫結至任何 C++ 可呼叫項目 (函式指標、lambda 等)，並具有相容的 operator() 簽章。建構的處理常式會解碼 XLA FFI 呼叫框架 (由穩定的 C API 定義)、檢查所有參數的類型，並將解碼結果轉送至使用者定義的回呼。
XLA FFI 繫結大量採用範本超程式設計，以便將建構的處理常式編譯為最有效率的機器碼。每個自訂呼叫參數的執行時間負擔，大約是幾奈秒。
以範本特化形式實作的 XLA FFI 自訂點，使用者可以定義如何解碼自訂型別，也就是可以為使用者定義的 enum class 型別定義自訂解碼。
從自訂呼叫傳回錯誤
自訂呼叫實作項目必須傳回 xla::ffi::Error 值，向 XLA 執行階段發出成功或錯誤信號。這與 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 需要等級 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 API 會從 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 支援自動解碼以 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]。(當然可以，而且應該這麼做！- 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);
在 GPU 上建立自訂呼叫
使用 XLA FFI 註冊 GPU 自訂呼叫幾乎相同，唯一不同的是，您需要要求基礎平台串流 (CUDA 或 ROCM 串流)，才能在裝置上啟動核心。以下是 CUDA 範例，可執行與上述 CPU 程式碼相同的運算 (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 自訂呼叫函式仍是在 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 解壓縮元組，再將其傳遞至自訂呼叫。
另一方面，元組 outputs 可讓您執行其他作業。
使用元組輸出的顯而易見原因，是自訂呼叫 (或任何其他 XLA 運算) 會透過元組輸出傳回多個獨立陣列。
但較不明顯的是，元組輸出也是提供自訂呼叫暫時記憶體的方式。可以，輸出可以代表暫時緩衝區。舉例來說，輸出緩衝區具有作業可寫入的屬性，且在寫入後可從中讀取。這正是您對暫時緩衝區的要求。
在上述範例中，假設我們想將 F32[1024] 用做暫時緩衝區。
接著，我們會像上述一樣編寫 HLO，但不會讀取自訂呼叫輸出內容的元組索引 1。