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

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);
    }
};

这是做事的正确方法吗？即理想情况下，在我的 REGISTER_OP 中使 cdata 成为 Input 或 Attr 会很好，但我认为这不会 link 正确地达到常量数据。我认为 cudaGetSymbolAddress 是必要的...
安全吗？即我是否会通过在提供的流中加入我自己的 cuda 命令和 memcpys 来干扰 tensorflow 的 GPU Stream Executor？
context->allocate_output 是获取一些固定内存的正确调用方法吗？查看 tensorflow 代码库表明存在临时和临时分配器，但我不知道它们是否暴露给用户...

编辑 1： 这会分配固定内存吗？（内存通常使用 cudaHostAlloc 分配，其页面被固定用于 DMA 传输到 GPU，即它们被阻止被 OS 换出）。

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

是的，cudaGetSymbolAddress 是必需的。常量内存是特定于内核的，不应该
不应该。只需确保流执行中的操作顺序正确并正确同步即可。
是输出是内核将作为操作结果写入的内存。暂存内存主要用于内核的单个操作所需的内存。一些 cudnn 内核，如卷积内核，使用它。见 tensorflow/kernels/conv_ops.cc

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

How does one transfer CUDA constant memory in tensorflow's C++ API

tensorflow