[Mlir-commits] [mlir] e7bb8dd - [mlir][linalg][bufferize] Relax rules for extract_slice/insert_slice matching

[Matthias Springer]

Author: Matthias Springer
Date: 2021-10-16T17:08:47+09:00
New Revision: e7bb8dd92983a918fbb370b135e664825f387580

URL: https://github.com/llvm/llvm-project/commit/e7bb8dd92983a918fbb370b135e664825f387580
DIFF: https://github.com/llvm/llvm-project/commit/e7bb8dd92983a918fbb370b135e664825f387580.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
    mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
index 21e0c1ce8cb8f..7e9a0e0ed38a6 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 c4173d6244d75..ffcabaf4cf54c 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
 //===----------------------------------------------------------------------===//