此页面由 Cloud Translation API 翻译。
Switch to English

平铺布局

图1

图1显示了如何使用2x2平铺将数组F32 [3,5]布置在内存中。具有此布局的形状写为F32 [3,5] {1,0:(2,2)},其中1,0与尺寸的物理顺序(布局中的minor_to_major字段)有关,而(2,2)冒号表示通过2x2瓦片平铺物理尺寸。

直观地布局图块以覆盖形状,然后在每个图块中,然后不进行平铺就对元素进行布局,如上例所示,该示例的右侧部分显示了内存中的布局,包括添加的白色填充元素为了获得完整的2x2切片,即使原始数组边界不均匀也是如此。

填充中的多余元素不需要包含任何特定值。

给定形状和图块的平铺线性指数公式

没有平铺,元素e =在阵列(例如N,E N-1,...,E 1)与数组边界当d =(d N,d N-1,...,d 1)(d1为最次要维度)按以下位置的次要顺序排列:

linear_index(e,d)
= linear_index((e n ,e n-1 ,...,e 1 ),(d n ,d n-1 ,...,d 1 ))
= e n d n-1 ... d 1 + e n-1 d n-2 ... d 1 + ... + e 1

为了简化本文中的符号表示,我们假设图块的维数与数组相同。在XLA的切片实施中,通过将初始的主要最大尺寸保持不变,并且仅将切片应用到最次要尺寸,将其概括为尺寸较小的切片,因此指定的切片提到了物理尺寸的后缀。形状被平铺。

当使用大小为(t n ,t n-1 ,...,t 1 )的切片时,将数组中具有索引(e n ,e n-1 ,...,e 1 )的元素映射到该元素在最终版式中的位置:

linear_index_with_tile(e,d,t)
= linear_index((⌊e/t⌋,e mod t),(⌈d/t⌉,t))(算术是元素方式,(a,b)是串联)
= linear_index((⌊eN / T N⌋,...,⌊e1 / T 1⌋,E N模T N,...,E 1模吨1),(⌈dN / T N⌉, ...,⌈d1 / T 1⌉,T N,T N-1,...,T 1))
= linear_index((⌊eN / T N⌋,...,⌊e1 / T 1⌋),(⌈dN / T N⌉,...,⌈d1 / T 1⌉))∙T N t n-1 ... t 1 + linear_index((e n mod t n ,...,e 1 mod t 1 ),(t n ,t n-1 ,...,t 1 ))

布局可以被认为是具有两个部分:(⌊eN / T N⌋,...,⌊e1 / T 1⌋),其对应于在尺寸的瓦片的阵列的瓦片索引(⌈dÑ / T N⌉,...,⌈d1 / T 1⌉),和(e N模T N,...,E 1模吨1),其对应于一个内瓦片的索引。所述细胞功能出现在⌈dI / T I⌉因为如果瓦片超限较大阵列的边界,填充插入如图1无论是瓷砖和瓦片内的元件被递归布置而不平铺。

对于图1中的示例,对于(1,1,0,1)的组合坐标矢量,元素(2,3)具有图块索引(1,1)和块内索引(0,1)。对于(2、3、2、2)的组合向量,图块索引的边界为(2,3),图块本身为(2,2)。然后,将具有逻辑形状的索引(2,3)的元素的带有图块的线性索引为

linear_index_with_tile((2,3),(3,5),(2,2))
= linear_index((1,1,0,1),(2,3,2,2))
= linear_index((1,1),(2,3))∙2∙2 + linear_index((0,1),(2,2))
=(1∙3 +1)∙2∙2 +(0∙2 +1)
= 17。

平铺为焊盘重塑移调

基于图块的布局操作如下:
考虑一个尺寸为(d n ,d n-1 ,...,d1)的数组(d1是最小尺寸)。当使用大小为(t n ,t n-1 ,...,t 1 )的拼贴进行布局时(t 1是最小尺寸),可以在下面用pad-reshape-transpose来描述该拼贴方式。

  1. 该阵列被填充到(⌈dN / T N⌉∙T N,...,⌈d1 / T 1⌉∙1)。
  2. 每个维度i被分成(⌈dI / T I⌉,T I),即该阵列被整形以
    (⌈dN / T N⌉,T N,...,⌈d1 / T 1⌉,T 1)。
    此重塑本身没有物理布局更改,因此此重塑是位广播。如果未明确考虑平铺,则此重塑可以表示具有与填充形状相同数量的元素的任何形状-这里的示例就是如何以这种方式表示平铺。
  3. 通过将t n ,...,t 1移至最小尺寸,同时保持它们的相对顺序来进行转置,这样从最主要到最小尺寸的尺寸顺序变为
    (⌈dN / T N⌉,...,⌈d1 / T 1⌉,T N,...,T 1)。

最终形状带有前缀
(⌈dN / T N⌉,...,⌈d1 / T 1⌉),其描述了在每一维的片的数目。数组中的一个元素(e n ,...,e 1 )以最终形状映射到该元素:
(⌊eN / T N⌋,...,⌊e0 /吨0⌋,E N模T N,...,E 1模吨1)。不难看出,元素的线性指数符合预期的上述公式。

重复平铺

通过重复应用XLA,平铺图块变得更加灵活。

图2

图2显示了大小为4x8的数组如何通过两层平铺(第一个2x4然后是2x1)进行平铺。我们将此重复拼贴表示为(2,4)(2,1)。每种颜色表示一个2x4的图块,每个红色边框表示一个2x1的图块。数字以平铺格式表示该元素在内存中的线性索引。此格式与TPU上BF16所使用的格式匹配,不同之处在于初始图块较大,即图块为(8,128)(2,1),其中第二次以2x1进行图块的目的是将两个16位值收集到一起以符合TPU架构的方式形成一个32位值。

请注意,第二个或更高版本的图块可以同时引用次要块内维度,后者只是重新排列图块内的数据,如在此示例中使用(8,128)(2,1),但也可以引用主交叉图块之前的拼贴的尺寸。

使用图块组合尺寸

XLA的切片也支持组合尺寸。例如,可以先将F32 [2,7,8,11,10] {4,3,2,1,0}中的尺寸合并到F32 [112,110] {1,0}中,然后再将其与(2,3 )。所使用的图块是(∗,∗,2,∗,3)。此处,图块中的星号表示采用该尺寸并将其与下一个较小的尺寸组合。多个相邻维度可以一起归纳为一个维度。包含的尺寸由该尺寸的瓷砖中的-1值表示,否则在尺寸中作为尺寸尺寸无效。

更准确地说,如果形状的尺寸i通过图块中的星号消除,那么在应用平铺的先前定义之前,该尺寸将从被平铺的形状和图块矢量中删除,而尺寸i-1是形状的数组边界从d i-1增加到d i d i-1 。对图块矢量中的每个星号重复此步骤。