[Mlir-commits] [mlir] e44f405 - [mlir][bufferization] Fix bug in findValueInReverseUseDefChain

[Matthias Springer]

Author: Matthias Springer
Date: 2023-02-16T16:12:31+01:00
New Revision: e44f405bb41b8acf8ef28c4e41bc06429e87f3c3

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

LOG: [mlir][bufferization] Fix bug in findValueInReverseUseDefChain

`alwaysIncludeLeaves` was not respected by all code paths.

Differential Revision: https://reviews.llvm.org/D144187

Added: 
    

Modified: 
    mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
    mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index f25a8eb0cf6c9..1bba60d06c57a 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -505,7 +505,8 @@ llvm::SetVector<Value> AnalysisState::findValueInReverseUseDefChain(
       if (followEquivalentOnly && a.relation != BufferRelation::Equivalent) {
         // Stop iterating if `followEquivalentOnly` is set but the alias is not
         // equivalent.
-        result.insert(value);
+        if (alwaysIncludeLeaves)
+          result.insert(value);
       } else {
         workingSet.insert(a.opOperand->get());
       }

diff  --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination.mlir
index 753840572f4b3..aa6e6a1dbe051 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination.mlir
@@ -204,3 +204,20 @@ func.func @eleminate_multiple_ops(%t: tensor<?xf32> {bufferization.buffer_layout
   %r1 = tensor.insert_slice %if into %t[42][%sz][1]: tensor<?xf32> into tensor<?xf32>
   return %r1: tensor<?xf32>
 }
+
+// -----
+
+// This is a regression test. Make sure that the tensor.extract_slice is not
+// eliminated.
+
+// CHECK-LABEL: func.func @regression_do_not_eliminate_non_empty(
+//       CHECK:   memref.subview
+//       CHECK:   memref.subview
+//       CHECK:   memref.copy
+func.func @regression_do_not_eliminate_non_empty(
+    %t: tensor<10xf32>, %t2: tensor<10xf32>) -> tensor<10xf32> {
+  %1 = tensor.extract_slice %t[0] [5] [1] : tensor<10xf32> to tensor<5xf32>
+  %2 = tensor.insert_slice %1 into %t2[1] [5] [1]
+      : tensor<5xf32> into tensor<10xf32>
+  return %2 : tensor<10xf32>
+}
\ No newline at end of file