From e7bb8dd92983a918fbb370b135e664825f387580 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm@google.com>
Date: Sat, 16 Oct 2021 17:08:00 +0900
Subject: [PATCH] [mlir][linalg][bufferize] Relax rules for
 extract_slice/insert_slice matching

The rules were too restrictive, causing out-of-place bufferization when the result of two ExtractSliceOp is fed into an InsertSliceOp.

Differential Revision: https://reviews.llvm.org/D111861
---
 .../Linalg/Transforms/ComprehensiveBufferize.cpp   |  3 ---
 .../comprehensive-module-bufferize-analysis.mlir   | 26 ++++++++++++++++++++++
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
index 21e0c1c..7e9a0e0 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
@@ -1193,9 +1193,6 @@ bool BufferizationAliasInfo::areEquivalentExtractSliceOps(
     return false;
   if (!sameOffsetsSizesAndStrides(st, sti, isEqualConstantIntOrValue))
     return false;
-  // TODO: Is the following needed?
-  if (!equivalentInfo.isEquivalent(st.result(), sti.source()))
-    return false;
   return true;
 }
 
diff --git a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir
index c4173d6..ffcabaf 100644
--- a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir
+++ b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir
@@ -278,6 +278,32 @@ func @extract_slice_to_linalg_write_use(
   return %D, %E: tensor<4x4xf32>, tensor<4x4xf32>
 }
 
+// -----
+
+// CHECK-LABEL: func @insert_slice_double_extract_slice
+func @insert_slice_double_extract_slice(
+    %s1: index, %s2: index, %s3: index, %s4: index, %A: tensor<8x6xf32>,
+    %B: tensor<6x6xf32>, %C: tensor<30x20xf32> {linalg.inplaceable = true})
+        -> tensor<30x20xf32> {
+  //      CHECK: tensor.extract_slice
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  %15 = tensor.extract_slice %C[%s3, %s4] [%s1, %s2] [1, 1] : tensor<30x20xf32> to tensor<?x?xf32>
+
+  //      CHECK: linalg.matmul
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  %18 = linalg.matmul ins(%A, %B : tensor<8x6xf32>, tensor<6x6xf32>) outs(%15 : tensor<?x?xf32>) -> tensor<?x?xf32>
+
+  //      CHECK: tensor.extract_slice
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  %19 = tensor.extract_slice %18[0, 0] [%s1, %s2] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+
+  //      CHECK: tensor.insert_slice
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  %20 = tensor.insert_slice %19 into %C[%s3, %s4] [%s1, %s2] [1, 1] : tensor<?x?xf32> into tensor<30x20xf32>
+
+  return %20 : tensor<30x20xf32>
+}
+
 //===----------------------------------------------------------------------===//
 // Transitive cases
 //===----------------------------------------------------------------------===//
-- 
2.7.4