PyTorch - 内核在哪里启动？

Question

PyTorch - 内核在哪里启动？

3

我需要获取有关PyTorch启动的内核的信息。例如，调用堆栈信息，如“main.py:24 -> ... -> callkernel.py:53”将是有益的。是否有任何方法可以在PyTorch应用程序执行期间收集此信息？我目前正在搜索PyTorch的源代码，但仍然找不到启动CUDA内核的行。我的问题有两个：

我能否在内核启动时获取调用堆栈？
有人能向我展示PyTorch源代码中内核启动的示例吗？

- Serkan Göktaş

1个回答

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Kurt Mohler · Accepted Answer

要获得有用的堆栈跟踪，您很可能需要使用调试符号构建pytorch（构建说明在这里）。我不确定是否有任何可供下载的调试版本。但是，如果没有一些背景知识，堆栈跟踪可能没有太多意义，因此这里是代码库中定义事物的一般概述：

PyTorch中的大多数运算符都作为C++ at::native命名空间函数在pytorch/aten/src/ATen/native中实现。当构建PyTorch时，代码生成脚本会自动生成native_functions.yaml中定义的运算符的Python函数和Python到C++的绑定，并且生成的代码没有检入存储库（因此如果您想要查看codegen中发生了什么，您必须读取脚本或者自行构建PyTorch）。

一个at::native操作符通常会调用该操作符的设备分发函数，该函数通常带有后缀_stub。分发函数会检查参数所在的设备（cpu、cuda等），然后运行特定于设备的实现。然后，另一个分发会发生，调用数据类型特定的实现。

举个例子，add.out操作符（当你在Python中执行torch.add(..., out=...)时调用）在这里声明。代码生成器生成了将Python函数绑定到at::native::add_out所需的一切内容，该函数在这里定义。请注意，该函数调用add_stub，这是设备分发函数。

一个CPU实现的add_stub已经在这里注册并且在这里作为add_kernel实现。一个CUDA实现已经在这里注册并且在这里作为add_kernel_cuda实现。请注意，这两个实现都使用了TensorIteratorBase对象。简而言之，这个对象将迭代张量输入中应该相加的每一对元素。

在 add_kernel 和 add_kernel_cuda 中还有另一个调度程序，它根据参数的数据类型选择不同的实现。这些不同的数据类型实现是从共享模板函数生成的。您可以看到 CPU 函数也有一个矢量化和非矢量化操作的不同实现，而 CUDA 实现只有这个。

如果您想查看完整的堆栈跟踪，可以使用 gdb --args python <script name> 运行脚本，并为您想要的特定内核创建断点。再次强调，需要调试符号才能理解它。