[mlir] Tighten the rules around folding TensorLoadOp

[nicolasvasilache]

Summary

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

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

 }
 
 //===----------------------------------------------------------------------===//
 // TensorLoadOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult TensorLoadOp::fold(ArrayRef<Attribute>) {
   if (auto tensorToMemref = memref().getDefiningOp<TensorToMemrefOp>())
+    // 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 {};
 }
 
 //===----------------------------------------------------------------------===//
 // TensorToMemrefOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult TensorToMemrefOp::fold(ArrayRef<Attribute>) {