自定义 AllReduce 算子实现思路

引言

PyTorch 的灵活性来自它的调度能力。离不开其强大的调度系统（Dispatcher）。Dispatcher 负责将一个算子调用路由到最合适的后端实现，并在 Autograd、设备、分布式等多个维度上进行组合调度。理解 Dispatcher，有助于我们更好地扩展自定义算子、定位性能问题，并为框架优化打下坚实基础。

为什么要理解 Dispatcher？

在实际工程中，我们常会遇到算子未命中、性能不达预期或自定义算子行为异常等问题。这些问题的根源往往与 DispatchKey 的分发逻辑息息相关。Dispatcher 不是一个黑箱，它是 PyTorch 算子分发、后端选择与扩展能力的关键枢纽。

核心概念速览

• Operator：算子在 IR 层的定义，包含 Schema 与元信息。
• DispatchKey：调度键，表示当前上下文可用的后端 / 变换策略。
• Kernel 注册：将具体实现与 DispatchKey 绑定，形成分发表。
• 后端与 Autograd 分层：不同后端与 Autograd 机制通过 DispatchKey 组合协同工作。

一个典型调用流程

dispatcher_call.cppC++

Tensor my_add(const Tensor& a, const Tensor& b) {
  auto op = Dispatcher::singleton()
    .findSchemaOrThrow("myops::add", "")
    .typed<Tensor (const Tensor&, const Tensor&)>();
  return op.call(a, b);
}

常见误区与排查建议

Dispatcher 关注点

小结

Dispatcher 是 PyTorch 核心能力的重要支柱，掌握其原理与最佳实践，不仅能帮助我们写出更高效的算子与后端实现，也是在系统层面理解 PyTorch 的重要一步。