2017-08-04 127 views
1

我正在为共享自定义功能代码的TensorFlow编写C++和CUDA操作系统。通常在CPU和CUDA实现之间共享代码时,如果编译CUDA,则会定义一个宏来将__device__说明符插入到函数签名中。在TensorFlow中是否有以这种方式共享代码的内置方式?TensorFlow CPU和CUDA代码共享

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

回答

0

事实证明,在TensorFlow中的下列宏将做我所描述的。

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