Torch Library API#
PyTorch C++ API 提供了扩展 PyTorch 核心算子库的能力,允许用户定义算子和数据类型。使用 Torch Library API 实现的扩展在 PyTorch 的 eager API 和 TorchScript 中均可使用。
有关 library API 的教程式介绍,请查阅使用自定义 C++ 算子扩展 TorchScript 教程。
宏#
-
TORCH_LIBRARY(ns, m)
用于定义一个函数,该函数将在静态初始化时运行,以在命名空间
ns
(必须是有效的 C++ 标识符,不带引号)中定义一个算子库。当您想定义一组 PyTorch 中尚不存在的新的自定义算子时,请使用此宏。
示例用法
TORCH_LIBRARY(myops, m) { // m is a torch::Library; methods on it will define // operators in the myops namespace m.def("add", add_impl); }
参数
m
绑定到一个 torch::Library 对象,该对象用于注册算子。对于任何给定的命名空间,只能有一个 TORCH_LIBRARY()。
-
TORCH_LIBRARY_IMPL(ns, k, m)
用于定义一个函数,该函数将在静态初始化时运行,以在命名空间
ns
(必须是有效的 C++ 标识符,不带引号)中为调度键k
(必须是 c10::DispatchKey 的非限定枚举成员)定义算子覆盖。当您想在新的调度键上实现一组已存在的自定义算子时(例如,您想为已存在的算子提供 CUDA 实现),请使用此宏。一个常见的用法模式是使用 TORCH_LIBRARY() 定义所有要定义的新算子的 schema,然后使用多个 TORCH_LIBRARY_IMPL() 块为 CPU、CUDA 和 Autograd 提供算子的实现。
在某些情况下,您需要定义适用于所有命名空间(而不仅仅是一个命名空间)的东西(通常是 fallback)。在这种情况下,请使用保留命名空间 _,例如:
TORCH_LIBRARY_IMPL(_, XLA, m) { m.fallback(xla_fallback); }
示例用法
TORCH_LIBRARY_IMPL(myops, CPU, m) { // m is a torch::Library; methods on it will define // CPU implementations of operators in the myops namespace. // It is NOT valid to call torch::Library::def() // in this context. m.impl("add", add_cpu_impl); }
如果
add_cpu_impl
是一个重载函数,请使用static_cast
指定您想要的重载(通过提供完整的类型)。
类#
-
class Library
此对象提供了定义算子和在调度键处提供实现的能力的 API。
通常,torch::Library 不是直接分配的;而是由 TORCH_LIBRARY() 或 TORCH_LIBRARY_IMPL() 宏创建的。
torch::Library 上的大多数方法都返回自身的引用,支持方法链式调用。
// Examples: TORCH_LIBRARY(torchvision, m) { // m is a torch::Library m.def("roi_align", ...); ... } TORCH_LIBRARY_IMPL(aten, XLA, m) { // m is a torch::Library m.impl("add", ...); ... }
公有函数
-
inline Library &def(c10::FunctionSchema &&s, const std::vector<at::Tag> &tags = {}, _RegisterOrVerify rv = _RegisterOrVerify::REGISTER) &
声明一个带有 schema 的算子,但不为其提供任何实现。
您应随后使用 impl() 方法提供实现。所有模板参数都会被推断出来。
// Example: TORCH_LIBRARY(myops, m) { m.def("add(Tensor self, Tensor other) -> Tensor"); }
- 参数
raw_schema – 要定义的算子的 schema。通常,这是一个
const char*
字符串字面量,但 torch::schema() 接受的任何类型在此处都可接受。
-
inline Library &set_python_module(const char *pymodule, const char *context = "")
声明后续所有 def 的算子的 fake impls 可以在给定的 Python 模块 (pymodule) 中找到。
这会注册一些帮助文本,如果找不到 fake impl,就会使用这些文本。
参数
pymodule: Python 模块
context: 我们可能会在错误消息中包含此信息。
-
inline Library &impl_abstract_pystub(const char *pymodule, const char *context = "")
已弃用;请改用 set_python_module。
-
template<typename NameOrSchema, typename Func>
inline Library &def(NameOrSchema &&raw_name_or_schema, Func &&raw_f, const std::vector<at::Tag> &tags = {}) & 为一个 schema 定义一个算子,然后为其注册一个实现。
如果您不打算使用 dispatcher 来构建算子实现,通常会使用此方法。它大致等同于调用 def() 再调用 impl(),但如果您省略算子的 schema,我们将从 C++ 函数的类型中推断出来。所有模板参数都会被推断出来。
// Example: TORCH_LIBRARY(myops, m) { m.def("add", add_fn); }
- 参数
raw_name_or_schema – 要定义的算子的 schema,或者仅为算子名称(如果 schema 从
raw_f
推断)。通常是一个const char*
字面量。raw_f – 实现此算子的 C++ 函数。此处接受 torch::CppFunction 的任何有效构造函数;通常您会提供函数指针或 lambda。
-
template<typename Name, typename Func>
inline Library &impl(Name name, Func &&raw_f, _RegisterOrVerify rv = _RegisterOrVerify::REGISTER) & 注册一个算子的实现。
您可以在不同的调度键处为单个算子注册多个实现(参见 torch::dispatch())。实现必须有相应的声明(来自 def()),否则无效。如果您计划注册多个实现,在 def() 算子时不要提供函数实现。
// Example: TORCH_LIBRARY_IMPL(myops, CUDA, m) { m.impl("add", add_cuda); }
- 参数
name – 要实现的算子名称。此处不要提供 schema。
raw_f – 实现此算子的 C++ 函数。此处接受 torch::CppFunction 的任何有效构造函数;通常您会提供函数指针或 lambda。
-
template<typename Func>
inline Library &fallback(Func &&raw_f) & 为所有算子注册一个 fallback 实现,当某个算子没有可用的特定实现时,将使用此 fallback 实现。
Fallback 必须关联一个 DispatchKey;例如,只能从命名空间为
_
的 TORCH_LIBRARY_IMPL() 中调用此函数。// Example: TORCH_LIBRARY_IMPL(_, AutogradXLA, m) { // If there is not a kernel explicitly registered // for AutogradXLA, fallthrough to the next // available kernel m.fallback(torch::CppFunction::makeFallthrough()); } // See aten/src/ATen/core/dispatch/backend_fallback_test.cpp // for a full example of boxed fallback
- 参数
raw_f – 实现 fallback 的函数。未装箱的函数通常不能用作 fallback 函数,因为 fallback 函数必须对每个算子都有效(即使它们的类型签名不同)。典型的参数是 CppFunction::makeFallthrough() 或 CppFunction::makeFromBoxedFunction()
-
inline Library &def(c10::FunctionSchema &&s, const std::vector<at::Tag> &tags = {}, _RegisterOrVerify rv = _RegisterOrVerify::REGISTER) &
-
class CppFunction
表示实现算子的 C++ 函数。
大多数用户不会直接与此类交互,除非通过错误消息:此类定义的构造函数定义了您可以通过接口绑定的“函数”类事物的允许集合。
此类抹去了传入函数的类型,但通过为函数推断的 schema 持久地记录了类型。
公有函数
-
template<typename Func>
inline explicit CppFunction(Func *f, std::enable_if_t<c10::guts::is_function_type<Func>::value, std::nullptr_t> = nullptr) 此重载接受函数指针,例如
CppFunction(&add_impl)
-
template<typename FuncPtr>
inline explicit CppFunction(FuncPtr f, std::enable_if_t<c10::is_compile_time_function_pointer<FuncPtr>::value, std::nullptr_t> = nullptr) 此重载接受编译时函数指针,例如
CppFunction(TORCH_FN(add_impl))
-
template<typename Lambda>
inline explicit CppFunction(Lambda &&f, std::enable_if_t<c10::guts::is_functor<std::decay_t<Lambda>>::value, std::nullptr_t> = nullptr) 此重载接受 lambda,例如
CppFunction([](const Tensor& self) { ...
})
公有静态函数
-
static inline CppFunction makeFallthrough()
这会创建一个 fallthrough 函数。
Fallthrough 函数会立即重新调度到下一个可用的调度键,但其实现比手动编写的相同功能的函数更有效率。
-
template<c10::BoxedKernel::BoxedKernelFunction *func>
static inline CppFunction makeFromBoxedFunction() 从具有签名
void(const OperatorHandle&, Stack*)
的盒装内核函数(boxed kernel function)创建函数;也就是说,它们在盒装调用约定中接收参数堆栈,而不是在原生的 C++ 调用约定中接收。盒装函数通常仅用于通过 torch::Library::fallback() 注册后端回退。
-
template<class KernelFunctor>
static inline CppFunction makeFromBoxedFunctor(std::unique_ptr<KernelFunctor> kernelFunctor) 从定义了
operator()(const OperatorHandle&, DispatchKeySet, Stack*)
(从盒装调用约定接收参数)并继承自c10::OperatorKernel
的盒装内核 functor 创建函数。与 makeFromBoxedFunction 不同,通过这种方式注册的函数还可以携带额外的状态,这些状态由 functor 管理;如果你正在编写适配器以连接到其他实现(例如,与注册内核动态关联的 Python 可调用对象),这会很有用。
-
template<typename FuncPtr, std::enable_if_t<c10::guts::is_function_type<FuncPtr>::value, std::nullptr_t> = nullptr>
static inline CppFunction makeFromUnboxedFunction(FuncPtr *f) 从未盒装内核函数(unboxed kernel function)创建函数。
这通常用于注册常见的运算符。
-
template<typename FuncPtr, std::enable_if_t<c10::is_compile_time_function_pointer<FuncPtr>::value, std::nullptr_t> = nullptr>
static inline CppFunction makeFromUnboxedFunction(FuncPtr f) 从编译时未盒装内核函数指针创建函数。
这通常用于注册常见的运算符。编译时函数指针可以允许编译器优化(例如,内联)对其的调用。
-
template<typename Func>
Functions#
-
template<typename Func>
inline CppFunction dispatch(c10::DispatchKey k, Func &&raw_f)# 创建一个与特定调度键(dispatch key)关联的 torch::CppFunction。
标记有 c10::DispatchKey 的 torch::CppFunctions 只有在调度程序确定这个特定的 c10::DispatchKey 是应该被调度的键时才会被调用。
此函数通常不直接使用,而是推荐使用 TORCH_LIBRARY_IMPL(),后者会隐式地在其主体内的所有注册调用中设置 c10::DispatchKey。
-
template<typename Func>
inline CppFunction dispatch(c10::DeviceType type, Func &&raw_f)# 接受 c10::DeviceType 的 dispatch() 便利重载。
-
inline c10::FunctionSchema schema(const char *str, c10::AliasAnalysisKind k, bool allow_typevars = false)#
从字符串构造 c10::FunctionSchema,并显式指定 c10::AliasAnalysisKind。
通常,schema 直接作为字符串传入,但如果您需要指定自定义别名分析(alias analysis),可以用对此函数的调用替换该字符串。
// Default alias analysis (FROM_SCHEMA) m.def("def3(Tensor self) -> Tensor"); // Pure function alias analysis m.def(torch::schema("def3(Tensor self) -> Tensor", c10::AliasAnalysisKind::PURE_FUNCTION));
-
inline c10::FunctionSchema schema(const char *s, bool allow_typevars = false)#
函数 schema 可以直接从字符串字面量构造。