IREE项目中tensor.extract_slice与tensor.expand_shape操作优化分析

在IREE编译器优化过程中，我们经常会遇到需要优化张量操作序列的情况。本文将深入分析一个特定的优化模式——如何有效地将tensor.expand_shape操作通过tensor.extract_slice操作进行"冒泡"（bubble up）优化。

问题背景

在IREE的DispatchCreation阶段，存在一个名为BubbleExpandThroughExtract的优化模式。该模式的核心目标是将extract_slice -> expand_shape的操作序列转换为expand_shape -> extract_slice的顺序。这种转换的主要优势在于可以将extract_slice操作克隆到其消费者dispatch中，当extract_slice操作不连续时，可以避免产生缓慢的内存拷贝操作。

当前优化限制

当前的优化实现有以下两个主要限制条件：

1. extract_slice操作不能修改那些被expand_shape扩展的维度
2. extract_slice操作必须是完全静态的

这些限制导致了许多潜在的优化机会被错过。例如，考虑以下两种情况：

// 可优化的情况
util.func public @possible(%arg0 : tensor<4096xf16>) -> (tensor<32x64xf16>) {
  %extracted_slice = tensor.extract_slice %arg0[0] [2048] [1] : tensor<4096xf16> to tensor<2048xf16>
  %expanded_239 = tensor.expand_shape %extracted_slice[[0, 1]] output_shape [32, 64] : tensor<2048xf16> into tensor<32x64xf16>
  util.return %expanded_239 : tensor<32x64xf16>
}

// 不可优化的情况
util.func public @impossible(%arg0 : tensor<2049xf16>) -> (tensor<32x64xf16>) {
  %extracted_slice= tensor.extract_slice %arg0[0] [2048] [1] : tensor<2049xf16> to tensor<2048xf16>
  %expanded_239 = tensor.expand_shape %extracted_slice[[0, 1]] output_shape [32, 64] : tensor<2048xf16> into tensor<32x64xf16>
  util.return %expanded_239 : tensor<32x64xf16>
}

优化条件分析

经过深入分析，我们发现这类优化的关键条件是：除最后一个输出维度外，所有其他输出维度的乘积必须能够整除输入（提取前）的形状。具体来说：

• 在第一个例子中，32（第一个输出维度）能够整除4096（输入形状），因此优化是可行的
• 在第二个例子中，32不能整除2049，因此优化不可行

动态形状处理

在实际应用中，我们还需要考虑动态形状的情况。例如：

util.func public @main(%arg0 : tensor<?xf16>, %val : index) -> (tensor<32x?xf16>) {
  %extracted_slice = tensor.extract_slice %arg0[0] [%val] [1] : tensor<?xf16> to tensor<?xf16>
  %cst32 = arith.constant 32 : index
  %div = arith.divsi %val, %cst32 : index
  %expanded_239 = tensor.expand_shape %extracted_slice[[0, 1]] output_shape [32, %div] : tensor<?xf16> into tensor<32x?xf16>
  util.return %expanded_239 : tensor<32x?xf16>
}

对于动态形状，我们需要确保输入张量的动态维度能够被扩展后的第一个维度（这里是32）整除。这通常需要在编译时通过形状推导或约束求解来验证。

实际应用案例

在实际的模型编译中，我们可能会遇到更复杂的多维情况：

%extracted_slice_446 = tensor.extract_slice %expanded_444
  [0, 31, 0, 0, 0, 0] 
  [%9, 1, 1, 8, 32, 128] 
  [1, 1, 1, 1, 1, 1] 
  : tensor<?x32x2x8x32x128xf8E4M3FNUZ> to tensor<?x8x32x128xf8E4M3FNUZ>
%expanded_447 = tensor.expand_shape %extracted_slice_446 
  [[0], [1], [2], [3, 4]] 
  output_shape [%9, 8, 32, 2, 64] 
  : tensor<?x8x32x128xf8E4M3FNUZ> into tensor<?x8x32x2x64xf8E4M3FNUZ>

在这种情况下，优化需要考虑多个维度的交互关系，确保在维度重组后，切片操作仍然能够保持数学上的等价性。

优化实现建议

基于上述分析，我们可以提出以下优化实现策略：

1. 放宽静态限制，支持动态形状的优化
2. 引入维度整除性验证机制，确保数学等价性
3. 开发更通用的维度关系分析工具，处理多维情况
4. 在优化前进行充分的形状推导和约束验证

通过这些改进，可以显著提高IREE编译器在处理复杂张量操作序列时的优化能力，从而生成更高效的代码。