我在 TensorFlow r1.5 中使用 C++ 和 CUDA 编写的一个操作的一部分涉及对张量进行归约。我实现了简单的交错归约算法,如这里所述。然而,似乎整个缓冲区并未被完全归约。
块归约的实现如下
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; //为了允许模板化、动态共享内存,我们将smem设置为uchar并重新解释为模板类型。 extern __shared__ __align__(sizeof(T)) unsigned char buffReduce[]; __syncthreads(); //设置此线程的贡献。如果超出范围则为0。 T *reduce = reinterpret_cast<T*>(buffReduce); reduce[tIdx] = (idx >= len) ? 0.0 : buffer[idx]; __syncthreads(); //块归约。#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) { //获取CUDA流。 const cudaStream_t &stream = dev.stream(); //获取归约操作的启动配置。 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); //将缓冲区重置为已知值。 std::vector<T> knownValsHost(len, 1.0); cudaMemcpyAsync(buffer, &knownValsHost[0], len * sizeof(T), cudaMemcpyHostToDevice, stream); CUSAFE(cudaStreamSynchronize(stream)); //将输出重置为零。 std::vector<T> tmp(blocks, 0.0); cudaMemcpyAsync(blockVals, &tmp[0], blocks * sizeof(T), cudaMemcpyHostToDevice, stream); CUSAFE(cudaStreamSynchronize(stream)); //在GPU上进行归约。 blockReduceDevice<T><<<blocks, threads, shared, stream>>>(buffer, blockVals, len); CUSAFE(cudaPeekAtLastError()); CUSAFE(cudaStreamSynchronize(stream)); //在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)); //获取将缓冲区复制到CPU内存并进行归约的结果。 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)); //打印一些诊断输出。 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: 32768Num CUDA Blocks: 10Num CUDA Threads Per Block: 1024GPU Result: 10240CPU Result: 32768
可以看出,使用std::accumulate
进行的CPU归约按预期工作(因为len == resCPU
)。这让我认为CUDA内核并未完全执行,因为blocks * threads != len
。
TensorFlow文档这里指出,应使用tensorflow/core/util/cuda_kernel_helper.h
头文件获取CUDA内核启动配置,该头文件可在这里找到。
TensorFlow为什么会提供一个不执行适当数量线程的启动配置?
我手动设置启动配置参数时也得到了类似的结果。
回答:
TensorFlow为什么会提供一个不执行适当数量线程的启动配置?
我猜是因为TensorFlow期望它运行的内核符合你的内核不符合的设计原则。TensorFlow返回的执行参数会将线程数限制为在给定设备上理论上可以并发运行的最大线程数。有关完整详情,请参见这里。
你的任务是编写一个符合该设计模式的内核,基本上就是让每个线程能够处理多个输入数据点。实际上,这意味着你需要简单地修改你的内核,如下所示:
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 //为了允许模板化、动态共享内存,我们将smem设置为uchar并重新解释为模板类型。 extern __shared__ __align__(sizeof(T)) unsigned char buffReduce[]; // cargo cult : __syncthreads(); //设置此线程的贡献。如果超出范围则为0。 T *reduce = reinterpret_cast<T*>(buffReduce); T threadsum = T(0); for(; idx < len; idx += stride) threadsum += buffer[idx]; //将线程局部部分归约存储到共享内存中 reduce[tIdx] = threadsum; __syncthreads(); // 等
[警告:显然从未编译或运行,风险自负]
基本上,这种设计将使每个线程尝试迭代通过所需的尽可能多的输入数据点,以确保内存合并的方式处理所有输入数据。