LibTorch 稳定 ABI¶
本说明最终将包含更多关于如何使用 torch/csrc/stable 中 API 的详细信息。目前,它包含一个内部表示的表格。
自定义扩展中的类型:最终用户自定义库中使用的类型。
StableIValue 表示:将类型稳定转换为 ABI 稳定方式下用户模型与 libtorch.so 之间的联络。
libtorch 中的类型:在 libtorch.so(或任何与 libtorch 二进制锁定的代码)中使用的类型。
Schema 类型:模式描述的类型,我们将其视为 native_functions.yaml 中 ATen 算子和通过 TORCH_LIBRARY 或 torch.library 注册到调度器的用户定义自定义算子的事实来源。
自定义扩展中的类型 |
StableIValue 表示 |
libtorch 中的类型 |
Schema 类型 |
---|---|---|---|
std::optional<S> |
*reinterpret_cast<(StableIValue*)*>,指向递归定义的 StableIValue 的指针 |
std::optional<T> |
Type? |
std::nullopt |
*reinterpret_cast<nullptr_t*> |
IValue() |
None |
RAIIATH |
AtenTensorHandle 的 *reinterpret_cast<uint64_t*> |
at::Tensor |
Tensor |
int32_t |
*reinterpret_cast<uint64_t*> |
at::ScalarType |
ScalarType |
int32_t |
*reinterpret_cast<uint64_t*> |
at::Layout |
Layout |
int32_t |
*reinterpret_cast<uint64_t*> |
at::MemoryFormat |
MemoryFormat |
bool |
*reinterpret_cast<uint64_t*> |
bool |
bool |
int64_t |
*reinterpret_cast<uint64_t*> |
int64_t |
int |
double |
*reinterpret_cast<uint64_t*> |
double |
float |
? |
? |
c10::Device |
Device |
? |
? |
c10::Stream |
Stream |
? |
? |
c10::complex |
complex |
? |
? |
at::Scalar |
Scalar |
? |
? |
std::string/const char*/ivalue::ConstantString |
str |
? |
? |
at::Storage |
Storage |
? |
? |
at::Generator |
Generator |
? |
? |
c10::List<T> |
Type[] |
? |
? |
ivalue::Tuple<T> |
(Type, …) |
? |
? |
c10::SymInt |
SymInt |
? |
? |
c10::SymFloat |
SymFloat |
? |
? |
c10::SymBool |
SymBool |
? |
? |
at::QScheme |
QScheme |
我们有信心支持的类型是表中已完成行的类型。对于有限的一组用例,我们还隐式支持任何可表示在 64 位内的字面类型作为 StableIValues,因为默认的 reinterpret_cast 将会成功。即使在自定义扩展中没有标准定义的 device 表示,您也可以在自定义 kernel 中使用 StableIValue 抽象来处理 c10::Device 等类型。例如,自定义算子可以接受一个 StableIValue device 作为参数,并直接通过 aoti_torch_call_dispatcher 将其传递给一个 aten 算子。