[Mlir-commits] [mlir] [mlir][tensor] Add more tensor.extract_slice canonicalization (PR #212974)

[Tuomas Kärnä]

================
@@ -117,6 +117,36 @@ func.func @linalg_effects(
 
 // -----
 
+// CHECK-LABEL: func @fold_extract_slice_of_fill_of_empty
+//   CHECK-NOT:   tensor.extract_slice
+//       CHECK:   %[[EMPTY:.*]] = tensor.empty() : tensor<4096xf32>
+//       CHECK:   %[[FILL:.*]] = linalg.fill ins(%[[CST:.*]] : f32) outs(%[[EMPTY]] : tensor<4096xf32>) -> tensor<4096xf32>
+//       CHECK:   return %[[FILL]] : tensor<4096xf32>
+func.func @fold_extract_slice_of_fill_of_empty(%cst : f32) -> tensor<4096xf32> {
+  %empty = tensor.empty() : tensor<4096x1xf32>
+  %filled = linalg.fill ins(%cst : f32) outs(%empty : tensor<4096x1xf32>) -> tensor<4096x1xf32>
+  %slice = tensor.extract_slice %filled[0, 0] [4096, 1] [1, 1]
+    : tensor<4096x1xf32> to tensor<4096xf32>
+  return %slice : tensor<4096xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_extract_slice_of_fill_of_empty_2d
+//   CHECK-NOT:   tensor.extract_slice
+//       CHECK:   %[[EMPTY:.*]] = tensor.empty() : tensor<32x48xf32>
+//       CHECK:   %[[FILL:.*]] = linalg.fill ins(%[[CST:.*]] : f32) outs(%[[EMPTY]] : tensor<32x48xf32>) -> tensor<32x48xf32>
+//       CHECK:   return %[[FILL]] : tensor<32x48xf32>
+func.func @fold_extract_slice_of_fill_of_empty_2d(%cst : f32) -> tensor<32x48xf32> {
+  %empty = tensor.empty() : tensor<64x96xf32>
+  %filled = linalg.fill ins(%cst : f32) outs(%empty : tensor<64x96xf32>) -> tensor<64x96xf32>
+  %slice = tensor.extract_slice %filled[0, 0] [32, 48] [1, 1]
----------------
tkarna wrote:

Thanks. As all of the 3 patterns remove tensor.extract_slice ops, they could indeed potentially affect bufferization.

Pattern 1. expand_shape + extract_slice that folds to a no-op. This seems like a relatively safe canolicalization pattern, and it's not clear if there are use cases where bufferization would be affected.

Pattern 2. tensor.empty + rank-reducing extract_slice that extracts a full slice. This also seems like a relatively safe pattern but is probably more frequent that 1.

Pattern 3. tensor.empty + linalg.fill + tensor.extract_slice. This is a more generic pattern and could indeed cause issues in bufferization. One option is to harden it like pattern2 so that it only applies to rank-reducing full slices (this would still address our use case in Lighthouse).

That said, on my behalf it's fine to move all of these to opt-in patterns.

https://github.com/llvm/llvm-project/pull/212974