如何在 Tensorflow 的 C++ API 中传输 CUDA 常量内存

Question

假设我有一个CUDA GPU内核，用于使用恒定内存的自定义tensorlfow操作：

__constant__ int cdata[100];
__global__ void frobulate(float * data)
{
   int i = blockDim.x*blockIdx.x + threadIdx.x;
   float value = data[i];
   for(int j=0; j < 100; ++j) {
      value += cdata[i];
   }
}

然后，当在我的Frobulate自定义操作中实现Compute方法时

class Frobulate : public tensorflow::OpKernel
{
public:
    void Compute(OpKernelContext * context) override
    {
        ...
        // Get the current device
        const Device & device = context->eigen_device<Eigen::GpuDevice>();
        // Local, mutating version of constant data.
        // For illustration purposes only
        int local_data[100];
        // Reason about our local shape
        TensorShape local_shape(100);
        // Create a pointer to hold allocated output
        Tensor * pinned_ary_ptr = nullptr;
        // Allocate memory for the complex_phase,
        // I don't think allocate_output is correct here...
        // but we need pinned host memory for an async transfer
        OP_REQUIRES_OK(context, context->allocate_output(
          0, local_shape, &pinned_ary_ptr));
        for(int i=0; i<100; ++i)
           { pinned_ary_ptr[i] = local_data[i]; }
        // Get the symbol address of cdata and enqueue an
        // async transfer on the device's stream
        int * d_cdata_ptr;
        cudaGetSymbolAddress((void **)&d_cdata_ptr, &cdata);
        cudaMemcpyAsync(d_cdata_ptr, pinned_ary_ptr, sizeof(int)*100,
           cudaMemcpyHostToDevice, device.stream());
        // Call the kernel
        frobulate<<<grid, blocks, 0, device.stream()>>>(data);
    }
};

1. 这是正确的做事方式吗？即理想情况下，在我的REGISTER_OP中使cdata为Input或Attr是好的，但我认为这不会正确地连接到常量数据。我认为cudaGetSymbolAddress是必要的
2. 它安全吗？即，我会通过在提供的流上排队我自己的cuda命令和memcpy来干扰tensorflow的GPU流执行器吗
3. context->allocate_output是获取固定内存的正确方法吗？查看tensorflow代码库表明存在临时和临时分配器，但我不知道它们是否向用户公开

编辑1:这会分配固定内存吗？（内存通常使用cudaHostAlloc分配，其页面被固定以用于DMA传输到GPU，即防止它们被操作系统交换）。

tensorflow::AllocatorAttributes pinned_allocator;
pinned_allocator.set_on_host(true);
pinned_allocator.set_gpu_compatible(true);
// Allocate memory for the constant data
OP_REQUIRES_OK(context, context->allocate_temp(
DT_UINT8, cdata_shape, &cdata_tensor,
pinned_allocator));    

Answer 1

1. 是的，cudaGetSymbolAddress是必需的。恒定内存是特定于内核的，不应该是

2. 不应该。只需确保流执行中的操作顺序正确并正确同步即可。

3. Yes输出是内核将作为操作结果写入的内存。scratch内存主要用于内核的单个操作所需的内存。一些像卷积核一样的核心，使用它。参见tensorflow/kernels/conv_ops.cc