c++ - 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() {
        //
    }
}