在算子的实现以及优化过程中,Memory Format是一个比较重要的概念,本文会对Pytorch中出现的Memory Format进行梳理,以及在算子优化过程中,了解下Memory Format可以发挥哪些作用,后续还会了解一下常见的Pytorch算子对于Memory Format的影响。
逻辑顺序和物理顺序
在物理层面中,Pytorch的数据都是以一维连续的方式存储在物理内存空间中,存储的顺序永远是NCHW(NCDHW),而逻辑上如何读取这块内存空间可以分成几种常见的形式,如channels fist,channels last, block format等等。对于不同block的选择主要是关注性能,某些op对channels first友好,有些则对channels last更加友好,要根据op的语义来看
Pytorch中的Memory Format概览
torch.contiguous_format
: 该Memory Format下,Tensor会以稠密且无重叠的方式排列,其stride会由大到小进行排列,也被称为NCHW。torch.channels_last
: 该Memory Format下,Tensor会以稠密且无重叠的方式排列stride的规律如下:strides[0] > strides[2] > strides[3] > strides[1] == 1
, Tensor 维度为4,也被称为NHWC向量。torch.channels_last_3d
: 该Memory Format下,Tensor会以稠密且无重叠的方式排列stride的规律如下:strides[0] > strides[2] > strides[3] > strides[4] > strides[1] == 1
, Tensor 维度为4,也被称为NDHWC向量。torch.preserve_format
: 该Memory Format用于clone输入的向量来保证与输入向量一致的Memory Format,如果输入向量是稠密且无重叠的话,输出向量的stride会与输入向量相同,其他情况下输出的向量的stride会与torch.contiguous_format
相一致,即从大到小。Onednn block format
:以上四种为Torch 官方文档中描述的Memory Format,大家在native_functions.yaml文件中可以看到mkldnn开头的方法,这里一般用的是Onednn优化的特殊的format,叫做block format。
Channels_First与Channels_last Format
在以下篇幅中,我们将Channels_First简称为CF,Channels_Last简称为CL。CL的存储顺序即将通道作为最临近维度来存储。下图为一个N=2,C=3,H=2,W=2的CF排列方式。
下图为一个N=2,C=3,H=2,W=2的CL的排列方式。与上图不同的是,上图使用W作为连续增长的维度,而在下图中,Channel作为了连续增长的维度。
Memory format之间的转换
# 创建一个Contiguous的Tensor
import torch
N, C, H, W = 10, 3, 32, 32
x = torch.empty(N, C, H, W)
print(x.stride()) # Outputs: (3072, 1024, 32, 1)
# 通过Tensor.to指定Memory Format来进行转换,也是Pytorch推荐的用法
x = x.to(memory_format=torch.channels_last)
print(x.shape) # Outputs: (10, 3, 32, 32) 维度信息是不会发生改变的
print(x.stride()) # Outputs: (3072, 1, 96, 3)
# 也可以使用如下方法,该方法会返回一个与输入数据相同的,指定Memory Format,需要注意的是,当输入就是该Memory Format的时候,就直接返回该对象
x = x.contiguous(memory_format=torch.channels_last)
print(x.stride()) # Outputs: (3072, 1, 96, 3)
#可以通过is_contiguous()来判断Memory 是否是指定类型的Memory Format,此外还要求该Tensor在内存中是连续的
print(x.is_contiguous(memory_format=torch.channels_last))
One DNN Memory Format
相比于上述常见的几个Memory Format,onednn的格式是追求更加极致的性能来进行优化的一种Format,后文称之为block format,之前常见的CF,CL,称之为plain format。block format与plain format相比能够实现更好的cache reuse 和 vectorization。下面以官方文档为例,来介绍下Onednn block format。
上图为一个普通的CF format,当将该逻辑地址映射为内存中的实际地址,其计算方法如下:
value(n, c, h, w) = n * CHW + c * HW + h * W + w
上图为一个典型的block format,之所以称之为block,是因为该Format将部分维度以指定的Size分成若干个block,上图是最为常见的nChw8c的layout,该layout对clannels这个维度进行了分block,size为8。其offset的计算方式如下:
offset_nChw8c(n, c, h, w) = n * CHW
+ (c / 8) * HW*8
+ h * W*8
+ w * 8
+ (c % 8)
他的排列方式就是每8个channel维度是相邻的,其实可以理解成NHWC,但是前面N每8个一组进行排列,还有常见的在AVX512平台的block format为nChw16c,估计原因是对于512,刚好一次能取16个32位的数据,有对该类型该兴趣的同学可以参考这个论文Distributed Deep Learning Using Synchronous Stochastic Gradient Descent。
小结
在这一节中,我们了解了Pytorch常见的Memory Format以及Memory Format转换之间的一些操作,还有ONEDNN库特有的blockFormat,在后续的章节中将探究以下不同类型的Format对性能的影响(感觉都是为了提高vectorization和cache reuse来提高性能)。
参考联接
CHANNELS LAST MEMORY FORMAT IN PYTORCH