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

[Adam Siemieniuk]

================
@@ -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]
----------------
adam-smnk wrote:

I agree in principle but I'd first evaluate scope of changes needed and tradeoffs in either case.
Even if it's only bufferization problem, it should be addressed before causing performance degradation in case these are severe.

All in all, +1 to gathering more data first.

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