我正在为具有共享自定义函数代码的 TensorFlow 使用 C++ 和 CUDA 编写一个 Op。通常当 CPU 和 CUDA 实现之间的代码共享时,如果为 CUDA 编译,人们会定义一个宏来将 __device__
说明符插入到函数签名中。在 TensorFlow 中是否有以这种方式共享代码的内置方式?
如何定义可以在 CPU 和 GPU 上运行的实用函数(通常是内联的)?
最佳答案
事实证明,TensorFlow 中的以下宏将执行我描述的操作。
namespace tensorflow{
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
void foo() {
//
}
}
关于c++ - TensorFlow CPU 和 CUDA 代码共享,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/45507534/