Torch 库 API#

PyTorch C++ API 提供了扩展 PyTorch 核心算子库以支持用户自定义算子和数据类型的能力。使用 Torch Library API 实现的扩展可以同时用于 PyTorch 的 eager API 和 TorchScript。

要了解该库 API 的教程风格介绍，请参阅使用自定义 C++ 算子扩展 TorchScript 教程。

宏#

TORCH_LIBRARY(ns, m) static void TORCH_LIBRARY_init_##ns(torch::Library &);                        \ static const torch::detail::TorchLibraryInit TORCH_LIBRARY_static_init_##ns( torch::Library::DEF,                                                     \ &TORCH_LIBRARY_init_##ns,                                                \ #ns,                                                                     \ std::nullopt,                                                            \ __FILE__,                                                               \ __LINE__);                                                               \ void TORCH_LIBRARY_init_##ns( torch::Library& 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) _TORCH_LIBRARY_IMPL(ns, k, m, C10_UID)

用于定义一个将在静态初始化时运行的函数，以在命名空间 ns（必须是有效的 C++ 标识符，不带引号）中为调度键 k（必须是 c10::DispatchKey 的未限定枚举成员）定义算子重载。

当您想为预先存在的自定义算子集提供新的调度键实现时（例如，您想为已存在的算子提供 CUDA 实现），请使用此宏。一种常见的用法是使用 TORCH_LIBRARY() 来定义您想定义的所有新算子的 schema，然后使用多个 TORCH_LIBRARY_IMPL() 块来为 CPU、CUDA 和 Autograd 提供算子实现。

在某些情况下，您需要定义适用于所有命名空间（而不仅仅是单个命名空间）的内容（通常是回退）。在这种情况下，请使用保留的命名空间 _，例如：

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

公共函数

Library(const Library&) = delete

Library &operator=(const Library&) = delete

Library(Library&&) = default

Library &operator=(Library&&) = default

~Library() = default

inline Library &def(c10::FunctionSchema &&s, const std::vector<at::Tag> &tags = {}, _RegisterOrVerify rv = _RegisterOrVerify::REGISTER) &

声明一个算子，但不提供任何实现。

您需要稍后使用 impl() 方法提供实现。所有模板参数都将自动推断。

// Example:
TORCH_LIBRARY(myops, m) {
  m.def("add(Tensor self, Tensor other) -> Tensor");
}

参数: raw_schema – 要定义的算子的 schema。通常是一个 const char* 字符串字面量，但 torch::schema() 接受的任何类型都可以在这里使用。

inline Library &def(const char *raw_schema, const std::vector<at::Tag> &tags = {}, _RegisterOrVerify rv = _RegisterOrVerify::REGISTER) &

inline Library &set_python_module(const char *pymodule, const char *context = "")

声明对于后续定义的所有算子，其伪实现（fake impls）可以在给定的 Python 模块 (pymodule) 中找到。

这会注册一些帮助文本，在找不到伪实现时会用到。

参数 (Args)

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 = {}) &

定义一个算子，然后为其注册一个实现。

这通常是您在不打算利用调度器来组织算子实现时使用的。它大致等同于调用 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。

c10::OperatorName _resolve(const char *name) const

template<typename Name, typename Func> inline Library &impl_UNBOXED(Name, Func*) &

inline Library &def(detail::SelectiveStr<false>, const std::vector<at::Tag> &tags[[maybe_unused]] = {}) &

inline Library &def(detail::SelectiveStr<true> raw_schema, const std::vector<at::Tag> &tags = {}) &

template<typename Func> inline Library &def(detail::SelectiveStr<false>, Func&&, const std::vector<at::Tag> &tags[[maybe_unused]] = {}) &

template<typename Func> inline Library &def(detail::SelectiveStr<true> raw_name_or_schema, Func &&raw_f, const std::vector<at::Tag> &tags = {}) &

template<typename Func> inline Library &impl(detail::SelectiveStr<false>, Func&&) &

template<typename Dispatch, typename Func> inline Library &impl(detail::SelectiveStr<false>, Dispatch&&key, Func&&raw_f) &

template<typename Func> inline Library &impl_UNBOXED(detail::SelectiveStr<false>, Func*) &

template<typename Func> inline Library &impl(detail::SelectiveStr<true> name, Func &&raw_f) &

template<typename Dispatch, typename Func> inline Library &impl(detail::SelectiveStr<true> name, Dispatch &&key, Func &&raw_f) &

template<typename Func> inline Library &impl_UNBOXED(detail::SelectiveStr<true>, Func*) &

template<typename Func> inline Library &fallback(Func &&raw_f) &

注册一个算子的回退实现，当没有特定算子实现可用时将被使用。

回退必须关联一个调度键；例如，只能从命名空间为 _ 的 TORCH_LIBRARY_IMPL() 调用此函数；例如，只有当命名空间为 _ 的 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 – 实现回退的函数。未打包的函数（Unboxed functions）通常不能用作回退函数，因为回退函数必须适用于每个算子（即使它们具有不同的类型签名）。常见的参数是 CppFunction::makeFallthrough() 或 CppFunction::makeFromBoxedFunction()。

template<class CurClass> inline torch::class_<CurClass> class_(const std::string &className)

template<class CurClass> inline torch::class_<CurClass> class_(detail::SelectiveStr<true> className)

template<class CurClass> inline detail::ClassNotSelected class_(detail::SelectiveStr<false> className)

void reset()

template<class CurClass> inline class_<CurClass> class_(const std::string &className)

template<class CurClass> inline class_<CurClass> class_(detail::SelectiveStr<true> className)

友元

friend class detail::TorchLibraryInit

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) { ... 。

})

~CppFunction()

CppFunction(const CppFunction&) = delete

CppFunction &operator=(const CppFunction&) = delete

CppFunction(CppFunction&&) noexcept = default

CppFunction &operator=(CppFunction&&) = default

inline CppFunction &&debug(std::string d) &&

公共静态函数

static inline CppFunction makeFallthrough()

创建回退函数。

回退函数会立即重新调度到下一个可用的调度键，但其实现比手动编写的相同功能的函数更有效率。

template<c10::BoxedKernel::BoxedKernelFunction *func> static inline CppFunction makeFromBoxedFunction()

从签名 void(const OperatorHandle&, Stack*) 的打包内核函数（boxed kernel function）创建函数；即，它们以打包的调用约定接收参数，而不是以原生的 C++ 调用约定。

打包函数通常只用于通过 torch::Library::fallback() 注册后端回退。

template<c10::BoxedKernel::BoxedKernelFunction_withDispatchKeys *func> static inline CppFunction makeFromBoxedFunction()

template<class KernelFunctor> static inline CppFunction makeFromBoxedFunctor(std::unique_ptr<KernelFunctor> kernelFunctor)

从打包的内核函数（boxed kernel functor）创建函数，该函数定义了 operator()(const OperatorHandle&, DispatchKeySet, Stack*)（以打包的调用约定接收参数），并继承自 c10::OperatorKernel。

与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)

从非装箱内核函数创建函数。

这通常用于注册常用运算符。

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> inline CppFunction dispatch(c10::DispatchKey k, Func &&raw_f)#

创建一个与特定 dispatch key 关联的 torch::CppFunction。

torch::CppFunctions 标记了 c10::DispatchKey，除非调度器确定应该调度此特定的 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 只需作为字符串传递，但如果您需要指定自定义别名分析，则可以用对此函数的调用替换字符串。

// 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 可以直接从字符串字面量构造。

Torch 库 API#

宏#

类#

函数#

文档

教程

资源