索引分析

本文档介绍了 HLO 索引分析，允许您以符号方式计算 HLO 操作的索引映射。索引映射是一个函数，用于将一个张量的索引映射到另一个张量的索引，例如 HLO 指令输出的索引和 HLO 指令输入的索引。

示例

发送从tensor<20xf32>到tensor<10x20x30xf32>的广播

p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}

从输出到输入的索引映射为 \((i, j, k) \mapsto (j)\) $i \in [0, 10]\(, \)j \in [0, 20]\( and \)k \in [0, 30]$。

设计初衷

XLA GPU 使用多种定制解决方案来推断聚合、运算数利用率和平铺方案（详见下文）。索引分析的目标是为此类用例提供可重复使用的组件。索引分析基于 MLIR 的仿射地图基础架构，并添加了 HLO 语义。

合并

当我们知道要读取输入的哪些元素/切片以计算输出的元素时，有关内存合并的推理变得可行。

操作数利用率

XLA 中的操作数利用率表示该指令的输出被完全使用的情况下使用的每个指令输入的使用量。目前，系统也不针对一般情况计算利用率。通过索引分析，您可以精确地计算利用率。

平铺

图块/切片是通过偏移量、大小和步长进行参数化的张量的超矩形子集。图块传播是一种使用操作本身的平铺参数来计算操作提供方/使用方的图块参数的方法。已有库可以为 softmax 和 dot 执行此操作。如果通过索引映射来表示图块传播，则可以让图块传播变得更加通用和稳健。

职能和领域

索引映射是一个函数， \(\boldsymbol{f}(\boldsymbol{d}, \boldsymbol{s})\)可将张量的多索引 \(\boldsymbol{d}\) 映射到 \(A\) 张量的元素/范围 \(B\)。参数 \(\boldsymbol{s}\) 是指存在于张量 \(B\)中，但不存在于张量中 \(A\)的维度索引范围。 \(A\).

例如，如果我们将 tensor<2x4x8x16xf32> 归约到 tensor<4x8xf32>，则从 2D 输出到 4D 输入的索引映射为\((d_0, d_1) \mapsto (s_0, d_0, d_1, s_1)\)，其中 \(d_i\) 是与输出张量索引对应的维度参数。参数 \(s_j\)会对多个值进行编码，也就是说，要计算输出的 \((d_0, d_1)\) 元素，我们需要 \((s_0, d_0, d_1, s_1)\) 输入元素，其中 \(s_0 \in [0, 2)\) 和\(s_1 \in [0, 16)\)。

此映射可以根据 HLO 指令的属性构造，或者可以组合未融合指令的映射来使索引进行融合。映射还具有一个域，用于指定映射存在哪些张量元素。

\[ \begin{eqnarray} \boldsymbol{f}(\boldsymbol{d}, \boldsymbol{s})\; &s.t.& \\ \boldsymbol{lb}_d &\leq& \boldsymbol{d} \leq \boldsymbol{ub}_d \\ \boldsymbol{lb}_s &\leq& \boldsymbol{s} \leq \boldsymbol{ub}_s \\ \boldsymbol{lb}_g &\leq& \boldsymbol{g}(\boldsymbol{d}, \boldsymbol{s}) \leq \boldsymbol{ub}_g \end{eqnarray} \]

由于我们希望尽可能减少重新计算，因此需要一个用于符号计算的库。XLA 已经依赖于 MLIR，因此我们使用 mlir::AffineMap，而不是编写符号算术库。

典型的 AffineMap 如下所示

(d0)[s0, s1] -> (s0 + 5, d0 * 2, s1 * 3 + 50)

AffineMap 有两种类型的参数，即维度和符号，可以分别用于 \(\boldsymbol d\) 和 \(\boldsymbol s\) 。AffineMap 不包含关于维度范围的任何元数据，因此我们必须自行提供此数据。

struct Range {
 int64_t lower_bound;
 int64_t upper_bound;
};

struct IndexingMap {
 mlir::AffineMap affine_map;
 std::vector<Range> dimension_ranges;
 std::vector<Range> symbol_ranges;
 llvm::DenseMap<mlir::AffineExpr, Range> expr_ranges;
};

dim_ranges 会对索引映射的维度参数 \(\boldsymbol{d}\) 的包容性框约束进行编码，这些约束条件通常与转置、减少、元素级、点等操作的输出张量的形状一致，但也存在一些例外情况，例如 HloConcatenateInstruction。

symbol_ranges 对参数可以采用的可能值进行编码。 \(\boldsymbol {s}\)

我们来举例研究，以便了解以上每个选项的实际含义。

将映射编入索引以执行未融合操作

按元素

对于元素级操作，索引映射是一个身份。

  p0 = f32[10, 20] parameter(0)
  p1 = f32[10, 20] parameter(1)
  add = f32[10, 20] add(p0, p1)

输出到输入的映射如下：

• 输出 -> input_0: \((d_0, d_1) \mapsto (d_0, d_1)\) for $\boldsymbol{d} \in [0,9] \times [0, 19]\(, i.e. \)\boldsymbol{d} \in {\rm Dom}（输出）$
• output -> input_1: \((d_0, d_1) \mapsto (d_0, d_1)\) for $\boldsymbol{d} \in {\rm Dom}（输出）$

输入到输出的映射

• input_i -> 输出： \((d_0, d_1) \mapsto (d_0, d_1)\) $\boldsymbol{d} \in {\rm Dom}(output)$

广播

广播意味着，当我们将输出映射到输入时，系统会移除一些维度；在将输入映射到输出时，系统会添加这些维度。

p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}

输出到输入的映射：

• 输出 -> 输入： \((d_0, d_1, d_2) \mapsto (d_1)\) $\boldsymbol{d} \in {\rm Dom}(output)$

输入到输出的映射

• 输入 -> 输出： \((d_0) \mapsto (s_0, d_1, s_1)\) $\boldsymbol{d} \in {\rm Dom}（输出）\( and \)\boldsymbol{s} \in [0, 9] \times [0, 29]$。

请注意，现在右侧显示了输入到输出映射的 \(\boldsymbol s\) 。这些是表示值范围的符号。例如，在此特定情况下，具有索引 \(d_0\) 的每个输入元素都会映射到 10x1x30 的输出切片。

Constant 和 Iota

方便的是，它们没有任何输入参数，因此没有可计算索引的内容。

Transpose

转置的索引映射是输入/输出维度的排列。

p0 = f32[3, 12288, 6, 128] parameter(0)
transpose = f32[3, 6, 128, 12288] transpose(p0), dimensions={0, 2, 3, 1}

输出到输入的映射：

• output -> input： \((d_0, d_1, d_2, d_3) \mapsto (d_0, d_3, d_1, d_2)\) 适用于\(\boldsymbol{d} \in {\rm Dom}(output)\)

输入到输出的映射如下：

• 输入 -> 输出： \((d_0, d_1, d_2, d_3) \mapsto (d_0, d_2, d_3, d_1)\) 对于\(\boldsymbol{d} \in {\rm Dom}(input)\)

反手扣篮

反向索引映射会将还原的维度更改为 $upper_bound(d_i) - d_i$：

p0 = f32[1, 17, 9, 9] parameter(0)
reverse = f32[1, 17, 9, 9] reverse(p0), dimensions={1, 2}

输出到输入的映射：

• 输出 -> 输入：$(d_0, d_1, d_2, d_3) \mapsto (d_0, -d_1 + 16, -d_2 + 8, d_3)\( for \)\boldsymbol{d} \in {\rm Dom}(output)$

输入到输出的映射如下：

• 输入 -> 输出：$(d_0, d_1, d_2, d_3) \mapsto (d_0, -d_1 + 16, -d_2 + 8, d_3)\( for \)\boldsymbol{d} \in {\rm Dom}(input)$

(Variadic)Reduce

变量归约具有若干输入和多个 init，从输出到输入的映射添加了缩减的维度。因此在某种意义上，它的行为类似于广播 与广播相反

p0 = f32[256,10] parameter(0)
p0_init = f32[] constant(-inf)
p1 = s32[256,10] parameter(1)
p1_init = s32[] constant(0)
reduce = (f32[10], s32[10]) reduce(p0, p1, p0_init, p1_init),
  dimensions={0}, to_apply=min

输出到输入的映射如下：

• 输出 -> input_j: \((d_0) \mapsto (s_0, d_0)\) for $\boldsymbol{d} \in {\rm Dom}(output)\( and \)\boldsymbol{s} \in [0, 9]$
• 输出 -> init_j: \((d_0) \mapsto ()\) for $\boldsymbol{d} \in {\rm Dom}(output)$

输入与输出的对应关系如下：

• input_i -> output_j： \((d_0, d_1) \mapsto (d_1)\) $\boldsymbol{d} \in {\rmDom}（输入）$
• init_i -> output_j： \(() \mapsto (s_0)\) 用于 \(\boldsymbol{s} \in [0, 9]\)

时长： \(i, j = 0, \ldots, INPUT\\_COUNT\)。

切片

切片从输出到输入的索引会形成分行索引映射，这对输出的每个元素都有效。从输入到输出的映射受限于输入中元素的步长范围。

p0 = f32[10, 20, 50] parameter(0)
slice = f32[5, 3, 25] slice(f32[10, 20, 50] p0),
  slice={[5:10:1], [3:20:7], [0:50:2]}

输出到输入的映射：

• output -> input： \((d_0, d_1, d_2) \mapsto (d_0 + 5, 7d_1 + 3, 2d_2)\) 适用于\(\boldsymbol{d} \in {\rm Dom}(output)\)

输入到输出的映射如下：

• 输入 -> 输出： \((d_0, d_1, d_2) \mapsto (d_0, d_1 / 7, d_2 / 2)\) 对于\(\boldsymbol{d} \in [5, 9] \times [3, 19] \times [0, 49]\) ，步长为 $[1, 7, 2]$。

TBD：从输入到输出的索引

调整形状

重塑有不同口味。

收起形状

这是一个从 N-D 到 1D 的“线性”重塑。

p0 = f32[4,8] parameter(0)
reshape = f32[32] reshape(p0)

输出到输入的映射：

• 输出 -> 输入： \((d_0) \mapsto (d_0 / 8, d_0 \mod 8)\) $\boldsymbol{d} \in {\rm Dom}(output)$

输入到输出的映射如下：

• 输入 -> 输出： \((d_0, d_1) \mapsto (8 d_0 + d_1)\) $\boldsymbol{d} \in {\rm Dom}(input)$。

展开形状

这是一个反向“折叠形状”操作，它将一维输入重塑为 N-D 输出。

p0 = f32[32] parameter(0)
reshape = f32[4, 8] reshape(p0)

输出到输入的映射：

• 输出 -> 输入： \((d_0, d_1) \mapsto (8 d_0 + d_1)\) $\boldsymbol{d} \in {\rm Dom}(output)$

输入到输出的映射如下：

• 输入 -> 输出： \((d_0) \mapsto (d_0 / 8, d_0 \mod 8)\) $\boldsymbol{d} \in {\rm Dom}(input)$。

常规调整

这些是重塑操作，无法表示为单个展开或收起形状。它们只能表示为由 2 个或更多个展开或收起形状组成的组合。

示例 1：线性化-去线性化。

p0 = f32[4,8] parameter(0)
reshape = f32[2, 4, 4] reshape(p0)

这种重塑可以表示为以下组合：先将收起形状从 tensor<4x8xf32> 变为 tensor<32xf32>，然后再将形状展开为 tensor<2x4x4xf32>。

输出到输入的映射：

• 输出 -> 输入：$(d_0, d_1, d_2) \mapsto (2d_0 + (4d_1 + d_2) / 8, 4d_1 + d_2) \mod 8)$

时长： \(\boldsymbol{d} \in {\rm Dom}(output)\)

输入到输出的映射如下：

• 输入 -> 输出：$(d_0, d_1) \mapsto ((8d_0 + d_1) / 16, ((8d_0 + d_1) \mod 16) / 4, d_1 \mod 4)$

时长： \(\boldsymbol{d} \in {\rm Dom}(input)\)。

示例 2：展开和收起的子形状

p0 = f32[4, 8, 12] parameter(0)
reshape = f32[32, 3, 4] reshape(p0)

这种重塑可以表示为两种重塑的组合。第一个函数将最外层的维度 tensor<4x8x12xf32> 收起为 tensor<32x12xf32>，第二个维度则将最内层的维度 tensor<32x12xf32> 展开为 tensor<32x3x4xf32>。

输出到输入的映射：

• output -> input： \((d_0, d_1, d_2) \mapsto (d_0 / 8, d_0 \mod 8, 4d_1 + d_2)\) for \(\boldsymbol{d} \in {\rm Dom}(output)\)

输入到输出的映射如下：

• input -> output： \((d_0, d_1, d_2) \mapsto (8d_0 + d_1, d_2 / 4, d_2 \mod 4)\) 适用于 \(\boldsymbol{d} \in {\rm Dom}(input)\)。

Bitcast

位播操作可以表示为转置-重形-转置序列。因此，其索引映射只是此序列的索引映射的组合。

串联

concat 的输出到输入映射为所有输入定义，但域不重叠，即一次仅使用其中一个输入。

p0 = f32[3,50] parameter(0)
p1 = f32[3,30] parameter(1)
concat = f32[3,80] concatenate(f32[3,50] p0, f32[3,30] p1),
  dimensions={1}

输出到输入的映射：

• output -> input 1：

\((d_0, d_1) \mapsto (d_0, d_1)\) - \(\boldsymbol{d} \in [0, 2] \times [0, 49]\)

• output -> input 2：

\((d_0, d_1) \mapsto (d_0, d_1 - 50)\) $\boldsymbol{d} \in [0, 2] \times [50,79]$

输入与输出的映射关系如下：

• 输入 1 -> 输出： \((d_0, d_1) \mapsto (d_0, d_1)\) $\boldsymbol{d} \in {\rm Dom}(input_1)$。
• 输入 2 -> 输出： \((d_0, d_1) \mapsto (d_0, d_1 + 50)\) $\boldsymbol{d} \in {\rm Dom}(input_2)$。

点（已实现输出到输入）

点的索引映射与缩减映射非常相似。

p0 = f32[4, 128, 256] parameter(0)
p1 = f32[4, 256, 64] parameter(1)
dot = f32[4, 128, 64] dot(p0, p1),
  lhs_batch_dims={0}, rhs_batch_dims={0},
  lhs_contracting_dims={2}, rhs_contracting_dims={1}

输出到输入的映射如下：

• output -> input_1: \((d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)\) （用于\(\boldsymbol{d} \in {\rm Dom}(output)\) 和 \(\boldsymbol{s} \in [0, 255]\)）
• output -> input_2: \((d_0, d_1, d_2) \mapsto (d_0, s_0, d_2)\) （用于\(\boldsymbol{d} \in {\rm Dom}(output)\) 和 \(\boldsymbol{s} \in [0, 255]\)）

输出的输入映射：

• input_1 -> output： \((d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)\) 用于\(\boldsymbol{d} \in {\rm Dom}(input_1)\) 和 \(\boldsymbol{s} \in [0, 63]\)
• input_2 -> output： \((d_0, d_1, d_2) \mapsto (d_0, s_0, d_1)\) 用于\(\boldsymbol{d} \in {\rm Dom}(input_2)\) 和 \(\boldsymbol{s} \in [0, 127]\)

缩短期限（待定）

垫子（待定）

为 Fusion 制作索引地图

融合操作的索引映射是指集群中的每个操作的索引映射。可能会出现以不同的访问模式多次读取某些输入的情况。

一个输入，多个索引映射

下面是一个 \(p_0 + p_0^T\)示例。

f {
  p0 = f32[1000, 1000] parameter(0)
  transpose_p0 = f32[1000, 1000]{0, 1} transpose(p0), dimensions={1, 0}
  ROOT a0 = f32[1000, 1000] add(p0, transpose_p0)
}

p0 的输出到输入索引映射将为 $(d_0, d_1) \mapsto (d_0, d_1)\( and \)(d_0, d_1) \mapsto (d_1, d_0)$。这意味着，要计算输出的一个元素，我们可能需要读取输入参数两次。

一个输入、去重索引映射

在某些情况下，索引映射实际上是相同的，即使并不明显。

f {
  p0 = f32[20, 10, 50] parameter(0)
  lhs_transpose_1 = f32[10, 20, 50] transpose(p0), dimensions={1, 0, 2}
  lhs_e = f32[10, 20, 50] exponential(lhs_transpose_1)
  lhs_transpose_2 = f32[10, 50, 20] transpose(lhs_e), dimensions={0, 2, 1}
  rhs_transpose_1 = f32[50, 10, 20] transpose(p0), dimensions={2, 1, 0}
  rhs_log = f32[50, 10, 20] exponential(rhs_transpose_1)
  rhs_transpose_2 = f32[10, 50, 20] transpose(rhs_log), dimensions={1, 0, 2}
  ROOT add = f32[10, 50, 20] add(lhs_transpose_2, rhs_transpose_2)
}

在这种情况下，p0 的输出到输入索引映射只是 $(d_0, d_1, d_2) \mapsto (d_2, d_0, d_1)$。

​

Softmax

softmax parameter 0 的输出到输入索引编制映射：

• \((d_0, d_1, d_2) \mapsto (d_0, d_1, d_2)\)
• \((d_0, d_1, d_2)[s_0] \mapsto (d_0, d_1, s_0)\)

for \(\boldsymbol{d} \in {\rm Dom}(output)\) 和 \(\boldsymbol{s} \in [0, 124]\)表示输入的最内层。

索引地图简化器

mlir::AffineMap 上游的默认简化器不能对维度/符号的范围做出任何假设。因此，它无法有效地简化包含 mod 和 div 的表达式。

我们可以利用仿射图中子表达式的下限和上限知识来进一步简化它们。

简化器可以重写以下表达式。

1. \((d_0, d_1) \mapsto (d_0 + d1 / 16, d1 \mod 16)\) $\boldsymbol{d} \in [0, 6] \times [0, 14]\( becomes \)(d_0, d_1) \mapsto (d_0, d_1)$
2. \( for \)\( becomes \)
3. $(d_0, d_1, d_2) \mapsto ((16d_0 + 4d_1 + d_2) /8, (16d_0 + 4d_1 + d_2) \mod 8)\( for \)d_i \in [0, 9] \mapsto ((16d_0 + 4d_1 + d_2) /8, (16d_0 + 4d_1 + d_2) \mod8)\( for \)d_i \in [0, 9] d_2 + d_2 + d_2\( becomes \)
4. \((d_0, d_1) \mapsto (-(-11d_0 - d_1 + 109) / 11 + 9)\) $\boldsymbol{d} \in [0, 9] \times [0, 10]\( becomes \)(d_0, d_1) \mapsto (d_0)$.

通过索引映射简化器，我们可以了解 HLO 中某些链式重塑会相互抵消。

p0 = f32[10, 10, 10] parameter(0)
reshape1 = f32[50, 20] reshape(p0)
reshape2 = f32[10, 10, 10] reshape(reshape1)

完成索引映射的构成及其简化之后，我们会得到

\((d_0, d_1, d_2) \mapsto (d_0, d_1, d_2)\).