矩阵转置的实现要点

# 矩阵转置的关键要素 来源：https://gudok.xyz/transpose/ 如前所述，分块算法在字长增大时效率更高。通用寄存器只有64位长，因此我们无法对其做更多优化。另一方面，x86_64拥有专用的SIMD寄存器和指令，其宽度为128、256或512位，具体取决于CPU世代。例如，使用256位寄存器，我们可以利用log2(32)=5级2×2块分解来构建高效的32×32转置器。可惜的是，使用SIMD扩展远比执行`+`或`>>`更具挑战性。大量高级指令使得我们可以用不同方式完成相同工作，且性能各异。此外，SIMD扩展的可用性严重依赖于CPU架构和世代。这里我们关注相对现代的x86_64 CPU世代。 首先简要回顾x64_64 CPU提供的SIMD扩展。x86 CPU最早广泛使用的SIMD扩展是MMX（"多媒体扩展"），如今已被遗忘。它引入了一组新寄存器，可以视为8/16/32位整数的向量，并提供了专用的指令集来并行操作这些寄存器的元素。随着时间的推移，越来越多的扩展被添加。新发布的扩展增加了寄存器宽度、寄存器数量，或引入了全新的指令。以下扩展按目标向量宽度分组，并按历史顺序列出： 我们将假设可以使用AVX2，但不包括更新的扩展。AVX2提供了256位寄存器和整数指令。这意味着我们将能够通过log2(32)=5级块分解来构建高效的32x32转置器。和之前一样，我们将逐层、逐行进行操作。每一行用一个`__m256i`类型的变量表示，我们将其视为`uint8_t[32]`类型的向量。该类型定义在`immintrin.h`头文件中，所有支持AVX2的主流编译器（Intel、GCC、Clang）都提供了该头文件。编译器可以将这种类型的变量自然地映射到256位AVX寄存器，就像`uint64_t`变量可以映射到通用64位寄存器一样。同一个头文件还提供了操作这些寄存器的内联函数。每个内联函数对应一条CPU指令，因此无需手动编写汇编代码。在我们的代码中，将仅依赖三条指令：`_mm256_shuffle_epi8`、`_mm256_blendv_epi8`、`_mm256_permute2x128_si256`。另外还会隐式使用两条指令：`_mm256_load_si256`和`_mm256_store_si256`，它们负责在内存和256位寄存器之间移动数据。我们无需手动调用它们，因为编译器在遇到`_mm256i r = *ptr;`和`*ptr = r;`时会自动插入相应指令。其他三条指令的语义更复杂，下面给出说明。 **洗牌（Shuffle）。** `_mm256_shuffle_epi8`（https://software.intel.com/sites/landingpage/IntrinsicsGuide/#text=_mm256_shuffle_epi8）指令接受一个源256位寄存器和一个256位控制寄存器，将输入视为8位值向量，根据控制寄存器对其进行洗牌并返回结果。控制寄存器为输出向量的每个索引（0..31）指定要从中复制元素的源向量索引，或指定一个特殊值表示目标位置必须置零。这种语义使得洗牌成为非常强大的指令。如果该指令没有限制，仅仅使用一条短指令就可以完成以下任何操作：相比之下，使用通用指令完成相同操作需要一长串掩码、移位和OR指令。然而，`_mm256_shuffle_epi8`有一个限制：元素只能在其所在的左、右128位*通道*内进行洗牌，而不能跨通道。例如，无法将元素从`src[30]`（左通道）复制到`dst[12]`（右通道）。因此，从技术上讲，`_mm256_shuffle_epi8`是分别应用于左、右128位通道的一对独立通用洗牌操作。在我们的算法中，我们将在前四个阶段（共五个）中使用洗牌指令来重新定位上一阶段矩阵行的寄存器中的元素。在第一层，我们交换相邻元素；在第二层，交换相邻的每两个元素组成的组；依此类推。下面的示意图展示了第三阶段洗牌的应用，此时我们交换相邻的每四个元素组成的组。 **置换（Permute）。** 在算法的第五阶段，我们转置2x2块矩阵，其中每个块为16x16，这意味着我们必须交换通道。AVX2中没有能够用一条指令完成跨通道通用洗牌的方法。但由于我们的情况非常特殊——只需交换通道——我们可以用另一条指令来实现： `_mm256_permute2x128_si256`（https://software.intel.com/sites/landingpage/IntrinsicsGuide/#text=_mm256_permute2x128_si256）。该指令接受两个输入数据寄存器`a`和`b`以及一个控制寄存器，允许构造一个寄存器，其两个128位通道各自可以从以下四个选项中选择：`a[0:127]`、`a[128:255]`、`b[0:127]`、`b[128:255]`。将控制寄存器设为0x01会使置换指令返回`a`并交换其通道；`b`在我们的情况中可以是任意值。 **混合（Blend）。** 混合（https://software.intel.com/sites/landingpage/IntrinsicsGuide/#text=_mm256_blendv_epi8）是另一个非常有用的指令。它接受两个源寄存器和一个掩码，并根据掩码从`src1[i]`或`src2[i]`中选择每个第`i`个元素来构造结果。总的来说，将洗牌/置换和混合相结合，我们可以按照以下方式从两个输入寄存器`a`和`b`中构建一个由几乎任意元素选择（记住跨通道限制）组成的值： 利用这条规则，我们准备构建32x32转置器。将会有log2(32)=5层。回想一下，在每一层，我们构建一个中间矩阵，其中每一行是应用于上一层矩阵两行计算结果。每个这样的计算可以用相同的操作序列表达。区别在于输入哪些行、应用哪些掩码，以及在最后一级我们需要使用置换而不是洗牌。下图展示了在第一阶段计算前两行的过程，此时我们转置32x32矩阵内部的每个2x2矩阵。原来的第a层变成了第b层，原来的第c层变成了第d层。在第二阶段，应用相同的操作序列，但现在处理的是元素对。在第五阶段，我们处理16个元素组成的组，并且使用置换代替洗牌。 从全局来看，完整的计算序列可以用下面的网络表达。每个顶点对应一个256位值，保存某个矩阵的整行（32个元素），同时还关联用于计算该值的指令。边表示依赖关系：必须先计算这些依赖值，才能计算目标值本身。第一列顶点加载`src`矩阵的行。接下来的五对列对应算法的五个阶段。每对列中的第一列是洗牌（或置换）的结果，第二列是混合的结果。因此，倒数第二列计算出的值包含转置矩阵的行，最后一列负责将它们保存到`dst`对应的内存中。远离其余网络的顶部顶点对应加载256位掩码，这些掩码需要传递给洗牌和混合。这些掩码应事先计算好。 为如此复杂的转置器编写代码并非易事。完全展开格式的手工编码（如Vec64）是不可能的，因为这需要编写超过400行高度重复的代码。使用嵌套循环（逐层、逐块、逐行）肯定是可行的，而且转置器会相当紧凑，但其性能将取决于编译器的选择。例如，编译器选择展开哪些循环以及展开到什么程度会影响执行单元的负载情况。第三种选择是借助某种脚本语言进行代码生成，这可能是创建中等大小内核（如本内核）最干净的方式。 为了利用代码生成，我们需要用顶点和边显式地实例化上述网络。每个顶点对应一段代码，执行单条向量指令并将结果写入唯一命名的变量。指令的参数是与左邻顶点关联的变量。这种方法的好处之一是我们可以试验顶点的不同排序。任何有效的拓扑排序都会产生相同的结果，但性能可能不同。计算可以按列进行，可以以n叉树的方式进行，可以根据随机生成的拓扑排序进行，或者以你能想到的任何巧妙方式进行。由于微架构细节，很难事先判断哪种排序在性能上更好。尽管如此，可以应用一些通用原则。一方面，我们希望生成的代码在任何步骤中都尽可能少地保留已计算但后续仍需使用的值。这是因为编译器将每个变量映射到某个寄存器，当寄存器用完时（AVX2只有16个寄存器），它必须执行*溢出*：将值保存到堆栈，并在需要时恢复。显然，溢出越少越好。另一个需要考虑的是，相同指令的调用应该分散开。连续发出32条相同的指令并非最佳选择，因为每种指令类型只有少数几个执行单元能处理它。如果连续发出大量相同指令，其他执行单元就会空闲。更好的做法是交错不同类型的指令，增加它们并行执行的机会。长话短说，经验上表现良好的解决方案是逐层执行所有计算，类似于标量向量化。每一步，上一层的两行矩阵作为输入，对它们应用四条指令：两次洗牌和两次混合，从而产生下一层矩阵的两行。这重复进行，一层接一层，一对行接一对行。按照这些规则计算值的顺序在上面的网络中由顶点附近的数字表示。 代码生成的32x32转置器必须插入到整个算法中，并再次精心选择预加载策略。由于我们使用代码生成，自然的选择是将所有16条预加载指令均匀地分布在生成的代码中。为什么是16？回忆一下，完成单个64x64块需要4次调用32x32转置器。因此，每次调用时我们需要预加载下一个块的16行。总共有393条指令，因此我们在生成的代码中每393/16=24条指令插入一次预加载。 在代码示例中，三个`origin`变量对应我们正在处理的32x32子矩阵的偏移量：源矩阵、目标矩阵和预取矩阵。`stride`值表示必须加到对应指针上的字节数，使得指针指向下一行的同一列。对于当前的转置器，两个`stride`都必须设置为`N`。 ` ` `（代码块从略，保留原文） 注意：由于原始消息中代码部分被截断，并且包含大量行号，我们保留其原样，仅翻译代码周围的文本描述。代码块本身未作翻译。