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

Question

我正在编写一个自定义 TensorFlow op，它需要能够读取 OpKernel Compute() 实现中张量的当前值，以便决定如何进行。 Eigen-unsupported TensorMap template (used by the tensorflow::TTypes typedefs) provides multiple overloads of the function call operator which I can use when the op is running on the host. In fact, the "Adding a New Op”教程示例代码使用函数调用运算符重载将值赋值到张量中。

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

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

我注意到 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().

的不同语义的代码

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

Device-agnostic way to read from a tensor?

c++

eigen3

tensorflow