形状对 GPU 上 tensorflow 内存分配的影响答案

【问题标题】：Effect of shape on memory allocation by tensorflow on GPU形状对 GPU 上 tensorflow 内存分配的影响
【发布时间】：2019-10-28 05:24:43
【问题描述】：

当我在具有 tensorflow 的 nvidia GPU 上创建形状为 (L, N, N) 的变量时，该变量是否占用了 L 个正方形块的平铺，其中每个块都是 N x N？

例如，三个正方形的记忆拼凑在一起将是：

# --- #
| NxN |
# --- #
| NxN |
# --- #
| NxN |
# --- #

tensorflow 将如何为形状为(a, b, N, N) 的变量分配 GPU 内存？

【问题讨论】：

【解决方案1】：

GPU 上的内存被分配为一维字节数组，如果可能的话进行对齐。在您的示例中，形状为 (L, N, N) 的张量将是大小为 8*L*N*N 字节的一维数组； (a, b, N, N) 大小写为 8*a*b*N*N 字节，假设为 float64 值。

您可以通过阅读 TensorFlow 源代码找到这一点：

gpu_cudamalloc_allocator.h 用void* AllocateRaw(size_t alignment, size_t num_bytes) 声明GPUcudaMallocAllocator 类，实现于
gpu_cudamalloc_allocator.cc。
allocator.h 声明 Allocator 类，GPUcudaMallocAllocator 的父类，它声明 void* AllocateRaw(size_t alignment, size_t num_bytes)。

由于这些 AllocateRaw 函数采用单一大小，所有 TensorFlow 内存都在底层分配在一维数组中。这在实践中很常见，因为将高阶数组表示映射到连续的一维内存块上很容易，并且它可以最大限度地提高设备上和设备之间传输的吞吐量。

【讨论】：

感谢您深入研究这么多细节。如果将两个矩阵分配为 GPU 上的一块内存，我们会更快地进行矩阵乘法吗？对于 CPU，我记得如果矩阵之一在乘法之前转置，矩阵乘法会更快。将 2D 矩阵存储为 1D 数组意味着沿着矩阵的行或列进行跨步内存访问。
好吧...所有的乘法可能同时发生在 GPU 上。所以我猜跨步访问并没有那么糟糕，因为所有访问实际上都是同时发生的。
GPU上矩阵乘法的最优算法与CPU算法不一样。我相信领导者是阻塞外积矩阵乘法 (ref)。使用共享块减少了跨步内存访问的影响，这种影响只发生在第一次读取时。