c++ - Tensorflow CUDA Reduction Op 没有完全减少

Question

我用 C++ 和 CUDA 编写的 TensorFlow r1.5 操作的一部分涉及对张量的缩减。我已经实现了简单的交错缩减算法，如所述here .但是，似乎并没有减少整个缓冲区。

block 减少的实现如下

template<typename T>
__global__
void blockReduceDevice(const T *buffer, T *out, size_t len) {
    const size_t tIdx = threadIdx.x;
    const size_t bIdx = blockIdx.x;
    const size_t bDim = blockDim.x;
    const size_t idx = bIdx * bDim + tIdx;

    //To allow templated, dynamic shared memory, we set the 
    //smem to be uchar and reinterpret as templated type.
    extern __shared__ __align__(sizeof(T)) unsigned char buffReduce[];
    __syncthreads();

    //Set contribution of this thread. 0 if out of bounds.
    T *reduce = reinterpret_cast<T*>(buffReduce);
    reduce[tIdx] = (idx >= len) ? 0.0 : buffer[idx];
    __syncthreads();

    //Block reduce.
#pragma unroll
    for (int i = bDim >> 1; i >= 1; i >>= 1) {
        if(tIdx < i) {
            reduce[tIdx] += reduce[tIdx + i];
        }
        __syncthreads();
    }

    if(tIdx == 0) {
        out[bIdx] = reduce[tIdx];
    }
}

上面的内核调用如下

template<typename T>
void testReduce(const T *buffer, T *blockVals, const GPUDevice &dev, size_t len) {
    //Get CUDA stream.
    const cudaStream_t &stream = dev.stream();

    //Get launch configuration for reduction operation.
    const auto reduceConfig = tensorflow::GetCudaLaunchConfig(len, dev);
    const size_t blocks = reduceConfig.block_count;
    const size_t threads = reduceConfig.thread_per_block;
    const size_t shared = threads * sizeof(T);


    //Reset buffer to known value.
    std::vector<T> knownValsHost(len, 1.0);
    cudaMemcpyAsync(buffer, &knownValsHost[0], len * sizeof(T), cudaMemcpyHostToDevice, stream);
    CUSAFE(cudaStreamSynchronize(stream));

    //Reset output to nought.
    std::vector<T> tmp(blocks, 0.0);
    cudaMemcpyAsync(blockVals, &tmp[0], blocks * sizeof(T), cudaMemcpyHostToDevice, stream);
    CUSAFE(cudaStreamSynchronize(stream));

    //Reduce on the GPU.
    blockReduceDevice<T><<<blocks, threads, shared, stream>>>(buffer, blockVals, len);
    CUSAFE(cudaPeekAtLastError());
    CUSAFE(cudaStreamSynchronize(stream));

    //Further reduce on the CPU.
    std::vector<T> blockValsHost(blocks, 0.0);
    cudaMemcpyAsync(&blockValsHost[0], blockVals, blocks * sizeof(T), cudaMemcpyDeviceToHost, stream);
    CUSAFE(cudaStreamSynchronize(stream));
    const T resGPU = std::accumulate(blockValsHost.begin(), blockValsHost.end(), static_cast<T>(0));

    //Get result when copying buffer to CPU memory and reducing.
    std::vector<T> bufferHost(len, 0.0);
    cudaMemcpyAsync(&bufferHost[0], buffer, len * sizeof(T), cudaMemcpyDeviceToHost, stream);
    CUSAFE(cudaStreamSynchronize(stream));
    const T resCPU = std::accumulate(bufferHost.begin(), bufferHost.end(), static_cast<T>(0));

    //Print some output for diagnostics.
    std::cout << "Length: " << len << std::endl;
    std::cout << "Num CUDA Blocks: " << blocks << std::endl;
    std::cout << "Num CUDA Threads Per Block: " << threads << std::endl;
    std::cout << "GPU Result: " << resGPU << std::endl;
    std::cout << "CPU Result: " << resCPU << std::endl;
}

在上面的测试用例中，给出了以下输出，其中所有缓冲区条目都设置为 1.0

Length: 32768
Num CUDA Blocks: 10
Num CUDA Threads Per Block: 1024
GPU Result: 10240
CPU Result: 32768

可以看出，使用 std::accumulate 的 CPU 减少按预期工作(如 len == resCPU)。这让我相信 CUDA 内核没有完全执行为 blocks * threads != len。

TensorFlow 文档指出 here CUDA 内核启动配置应该使用 tensorflow/core/util/cuda_kernel_helper.h header 获取，可以在 here 中找到.

出于什么原因，TensorFlow 会为我提供未执行适当线程数的启动配置？

我在手动设置启动配置参数时也收到了类似的结果。

Answer 1

For what reason would TensorFlow provide me with a launch configuration that does not execute the appropriate number of threads?

我猜是因为 Tensorflow 期望它运行的内核符合您的内核不符合的设计原则。 Tensorflow 返回的执行参数会将线程数限制为理论上可以在给定设备上运行的最大并发线程数。参见 here了解全部详情。

您的工作是编写一个符合该设计模式的内核，基本上是通过每个线程能够处理多个输入数据点。实际上，这意味着将您的内核修改成这样:

template<typename T>
__global__
void blockReduceDevice(const T *buffer, T *out, size_t len) {
    const size_t tIdx = threadIdx.x;
    const size_t bIdx = blockIdx.x;
    const size_t bDim = blockDim.x;
    const size_t idx = bIdx * bDim + tIdx;
    const size_t stride = gridDim.x * blockDim.x

    //To allow templated, dynamic shared memory, we set the 
    //smem to be uchar and reinterpret as templated type.
    extern __shared__ __align__(sizeof(T)) unsigned char buffReduce[];
    // cargo cult : __syncthreads();

    //Set contribution of this thread. 0 if out of bounds.
    T *reduce = reinterpret_cast<T*>(buffReduce);
    T threadsum = T(0);
    for(; idx < len; idx += stride)
        threadsum += buffer[idx];

    // store thread local partial reduction to shared memory
    reduce[tIdx] = threadsum;
    __syncthreads();

    // etc

[警告:显然从未编译或运行，使用风险自负]

基本上，此设计将使每个线程尝试遍历尽可能多的输入数据点，以确保内存合并的方式处理所有输入数据。