TensorFlow CPU 和 CUDA 代码共享

Question

我正在为共享自定义函数代码的 TensorFlow 编写 C++ 和 CUDA 中的 Op。通常在 CPU 和 CUDA 实现之间共享代码时，如果为 CUDA 编译，则会定义一个宏以将__device__说明符插入函数签名中。有没有一种内置的方法可以在TensorFlow中以这种方式共享代码？

如何定义可以在CPU和GPU上运行的实用程序函数(通常是内联的(？

Answer 1

事实证明，TensorFlow 中的以下宏将执行我所描述的操作。

namespace tensorflow{
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
void foo() {
//
}
}