在算子的实现以及优化过程中，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排列方式。 Channels First

下图为一个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来提高性能）。

参考联接

Pytorch 官方文档

CHANNELS LAST MEMORY FORMAT IN PYTORCH

Intel OneDNN 官方文档

Pytorch中 MemoryFormat（一）