[Mlir-commits] [mlir] f21d786 - [mlir] Tighten the rules around folding TensorLoadOp

[Nicolas Vasilache]

Author: Nicolas Vasilache
Date: 2021-03-04T17:48:09Z
New Revision: f21d78633a3d09d1bac904cf4b97f9f7e66b3094

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

LOG: [mlir] Tighten the rules around folding TensorLoadOp

`tensor_load(tensor_to_memref(x)) -> x` is an incorrect folding because it ignores potential aliasing.

This revision approximates no-aliasing by restricting the folding to occur only when tensor_to_memref
is immediately preceded by tensor_load in the same block. This is a conservative step back towards
correctness until better alias analysis becomes available.

Context: https://llvm.discourse.group/t/properly-using-bufferization-related-passes/2913/6

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

Added: 
    

Modified: 
    mlir/lib/Dialect/StandardOps/IR/Ops.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
index 7a71f09adf63..11a597ce4183 100644
--- a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
+++ b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
@@ -3899,7 +3899,11 @@ void SubTensorInsertOp::getCanonicalizationPatterns(
 
 OpFoldResult TensorLoadOp::fold(ArrayRef<Attribute>) {
   if (auto tensorToMemref = memref().getDefiningOp<TensorToMemrefOp>())
-    return tensorToMemref.tensor();
+    // Approximate alias analysis by conservatively folding only when no there
+    // is no interleaved operation.
+    if (tensorToMemref->getBlock() == this->getOperation()->getBlock() &&
+        tensorToMemref->getNextNode() == this->getOperation())
+      return tensorToMemref.tensor();
   return {};
 }