从PyTorch调用到torch_xla

You 2023-12-31 00:33:46  66481

分类专栏：资讯

从PyTorch调用到torch_xla

xla调用上面所说的宏进行注册的位置在RegisterXLA.cpp这个文件中（codegen的结果），如下：

ORCH_LIBRARY_IMPL(aten, XLA, m) {
  m.impl("abs",
  TORCH_FN(wrapper__abs));

  ...
}

其中，wrapper__abs的定义如下：

at::Tensor wrapper__abs(const at::Tensor & self) {
  return torch_xla::XLANativeFunctions::abs(self);
}

显然，这个定义和PyTorch框架内部的算子是完全一致的，只是修改了实现。而XLANativeFunctions::abs的实现可以在aten_xla_type.cpp中找到，如下所示：

at::Tensor XLANativeFunctions::abs(const at::Tensor& self) {
  XLA_FN_COUNTER("xla::");
  return bridge::AtenFromXlaTensor(XLATensor::abs(bridge::GetXlaTensor(self)));
}

到这里已经比较明朗了，注册之后，PyTorch上对于op的调用最终会进入torch_xla的native function中调用对应的op实现，而这些实现的根本都是对XLATensor进行操作，在最终操作执行完成之后，会将作为结果的XLATensor重新转换为torch Tensor，但要注意，这里的结果不一定被实际计算了，也可能只是记录了一下IR，将节点加入图中，这取决于具体的实现。

总结

其实torch-xla官方的文档里是有关于代码生成和算子注册这个过程的描述的，只不过一开始我没找到这个文档，走了一点弯路，但是自己探索也会觉得更明了这个过程。官方文档中的描述如下（节选）：

All file mentioned below lives under the xla/torch_xla/csrc folder, with the exception of codegen/xla_native_functions.yaml

xla_native_functions.yaml contains the list of all operators that are lowered. Each operator name must directly match a pytorch operator listed in native_functions.yaml. This file serves as the interface to adding new xla operators, and is an input to PyTorch's codegen machinery. It generates the below 3 files: XLANativeFunctions.h, RegisterXLA.cpp, and RegisterAutogradXLA.cpp
XLANativeFunctions.h and aten_xla_type.cpp are entry points of PyTorch to the pytorch_xla world, and contain the manually written lowerings to XLA for each operator. XLANativeFunctions.h is auto-generated through a combination of xla_native_functions.yaml and the PyTorch core native_functions.yaml file, and contains declarations for kernels that need to be defined in aten_xla_type.cpp. The kernels written here need to construct 'XLATensor' using the input at::Tensor and other parameters. The resulting XLATensor needs to be converted back to the at::Tensor before returning to the PyTorch world.
RegisterXLA.cpp and RegisterAutogradXLA.cpp are auto-generated files that register all lowerings to the PyTorch Dispatcher. They also include auto-generated wrapper implementations of out= and inplace operators.

大概意思就是实际上torch-xla就是根据xla_native_functions.yaml这个文件来生成算子的定义，然后再生成对应的RegisterXLA.cpp中的注册代码，这也跟PyTorch的codegen方式一致。

综合这一整个过程可以看出，PyTorch是保持了高度的可扩展性的，不需要多少侵入式的修改就可以将所有的算子全部替换成自己的，这样的方式也可以让开发者不用去关注dispatcher及其上层的实现，专注于算子本身的逻辑。

网站声明：如果转载，请联系本站管理员。否则一切后果自行承担。

本文链接：https://www.xckfsq.com/news/show.html?id=33071

赞同 0

评论 0 条

从PyTorch调用到torch_xla

从PyTorch调用到torch_xla

总结

相关文章

关注我们