c++ - 从张量读取设备不可知的方式？

Question

我正在编写一个自定义 TensorFlow op，它需要能够读取 OpKernel Compute() 实现中张量的当前值，以便决定如何进行。 Eigen 不受支持的 TensorMap template (由 tensorflow::TTypes typedefs 使用)提供函数调用运算符的多个重载，当 op 在主机上运行时我可以使用它们。事实上，“Adding a New Op”教程示例代码使用函数调用运算符重载将值赋给张量。

问题是当 op 在 GPU 上运行时，调用函数调用运算符重载会导致段错误。这是因为函数调用运算符重载了对 m_data 的取消引用。成员指针，但当张量驻留在 GPU 上时，这是一个设备指针。此外，与 this other question 不同，我认为我不能使用 HostMemory()，因为张量是临时的，由 GPU 代码使用。

我目前正在调用 cudaMemcpy() 将设备内存复制到主机，但我想使用与设备无关的方式从设备内存中读取——无论设备是 Eigen::ThreadPoolDevice , Eigen::GpuDevice ，或其他东西——因为我将能够使用相同的 Eigen 代码，但为每个支持的设备编译(“仿函数”方法)。

我注意到 TensorBase::eval() 方法，但我不确定这是否是我要找的。

Answer 1

每个 Eigen *Device 类型都有一个 memcpyDeviceToHost() 成员。 DefaultDevice 和 ThreadPoolDevice 的实现调用 memcpy()。 GpuDevice 的实现调用 cudaMemcpyAsync()。

一个问题是 SyclDevice 的 memcpyDeviceToHost() 具有不同的签名和不同的语义；而 DefaultDevice、ThreadPoolDevice 和 GpuDevice 具有:

void memcpyDeviceToHost(void* dst, const void* src, size_t n) const

.. 复制 n bytes，SyclDevice 有:

void memcpyDeviceToHost(void* dst, const T* src, size_t n) const

.. 复制 n T 个对象。

为避免潜在问题，最好将 src 设备指针静态转换为 const void*，这样当设备类型为是 SyclDevice，指出需要专门用于 SyclDevice::memcpyDeviceToHost() 的不同语义的代码。