如何插入 tensorflow 中使用的 cuda 运行时 API？

Question

我已经在带有 CUDA 10 和合适的 GPU 的 x86-64 Ubuntu 计算机上安装了 tensorflow（不是从源代码）。我的目标是拦截（使用 LD_PRELOAD）任何 tensorflow 应用程序的 CUDA 运行时 API。不幸的是，对于我的用例，我无法从源代码为不是 x86-64 的目标机器构建 tensorflow。我能够拦截我从用 C++ 编写的动态加载 cuda 运行时 API 的测试程序进行的 cudaLaunchKernel 调用，并且在第一次检查时我认为 python 应该类似地动态加载相同的 cuda .so - 我很困惑，因为 LD_PRELOAD 不是适用于在启用 cuda 的情况下运行的正常安装的 tensorflow 应用程序。我希望 tensorflow 中的 cudaLaunchKernel 调用应该被我的 LD_PRELOAD 插入库拦截。这个对吗？

Answer 1

看起来 Tensorflow 编写了 stream_executor 以避免必须使用 CUDA 的运行时 API，而是用开源代码包装 CUDA 的驱动程序 API (cuLaunchKernel) 本身。在此对 Tensorflow 的拉取请求中引用了这 一点，以允许插入被拒绝的 CUDA 的运行时 API。同样，在 TF 源代码（参见此处）中，我们看到正在积极使用 cu*/driver API 而不是运行时 API。