Question

我正在编写一个自定义TensorFlow op，它需要能够在OpKernel Compute（）实现中读取张量的当前值，以便决定如何继续。 Eigen-unsupported的recommended by composer（由m_data typedef使用）提供了函数调用操作符的多个重载，当op在主机上运行时我可以使用它。实际上，“TensorMap template”教程示例代码使用函数调用运算符重载来将值赋值给张量。

问题是当op在GPU上运行时，调用函数调用运算符会导致段错误。这是因为函数调用运算符重载取消引用Eigen::ThreadPoolDevice成员指针，但这是张量驻留在GPU上时的设备指针。此外，与Adding a New Op不同，我认为我不能使用HostMemory（），因为张量是临时的，并且由GPU代码使用。

我目前正在调用cudaMemcpy（）将设备内存复制到主机，但我想使用与设备无关的方式从设备内存中读取 - 设备是Eigen::GpuDevice，{{1}}，或其他东西 - 因为我可以使用相同的本征代码，但为每个支持的设备编译（“仿函数”方法）。

我注意到了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 字节，SyclDevice具有：

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

..复制n T个对象。

为避免潜在问题，最好将src设备指针静态强制转换为const void*，这样当设备类型为SyclDevice时，您将收到编译错误，精确定位需要专门用于SyclDevice :: memcpyDeviceToHost（）的不同语义的代码。

设备无关的从张量读取的方法？

1 个答案: