[mlir][bufferization] Do not canonicalize to_tensor(to_memref(x))

[springerm]

Summary

This is a partial revert of D128615.

to_memref(to_tensor(x)) always be folded to x. But to_tensor(to_memref(x)) cannot be folded in the general case because writes to the intermediary memref may go unnoticed.

mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp

if (auto toMemref = getMemref().getDefiningOp<ToMemrefOp>())
     // is no interleaved operation.
     if (toMemref->getBlock() == this->getOperation()->getBlock() &&
         toMemref->getNextNode() == this->getOperation())
       return toMemref.getTensor();
   return {};
 }
 
 namespace {
-/// Canonicalize bufferization.to_tensor + bufferization.to_memref.
-struct ToTensorToMemrefFolding : public OpRewritePattern<ToTensorOp> {
-  using OpRewritePattern<ToTensorOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(ToTensorOp toTensorOp,
-                                PatternRewriter &rewriter) const final {
-    auto toMemrefOp = toTensorOp.getMemref().getDefiningOp<ToMemrefOp>();
-    if (!toMemrefOp)
-      return failure();
-    rewriter.replaceOp(toTensorOp, toMemrefOp.getTensor());
-    return success();
-  }
-};
-
 struct DimOfToTensorFolder : public OpRewritePattern<tensor::DimOp> {
   using OpRewritePattern<tensor::DimOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(tensor::DimOp dimOp,
                                 PatternRewriter &rewriter) const override {
     auto memrefToTensorOp = dimOp.getSource().getDefiningOp<ToTensorOp>();
     if (!memrefToTensorOp)
       return failure();
 
     rewriter.replaceOpWithNewOp<memref::DimOp>(
         dimOp, memrefToTensorOp.getMemref(), dimOp.getIndex());
     return success();
   }
 };
 } // namespace
 
 void ToTensorOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
-  results.add<DimOfToTensorFolder, ToTensorToMemrefFolding>(context);
+  results.add<DimOfToTensorFolder>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // ToMemrefOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult ToMemrefOp::fold(ArrayRef<Attribute>) {
   if (auto memrefToTensor = getTensor().getDefiningOp<ToTensorOp>())

mlir/test/Dialect/SCF/canonicalize.mlir

%0:3 = scf.for %arg0 = %lb to %ub step %step iter_args(%arg1 = %t0, %arg2 = %t1, %arg3 = %t2)
 
     // CHECK-NEXT:   scf.yield %[[YIELD_T]] : tensor<128x128xf32>
     scf.yield %1, %2, %3 : tensor<128x128xf32>, tensor<128x128xf32>, tensor<128x128xf32>
 
   // CHECK-NEXT: }
   }
 
   // CHECK-NEXT: %[[R0:.*]] = bufferization.to_tensor %[[M0]] : memref<128x128xf32>
-  // CHECK-NEXT: return %[[R0]], %[[T1]], %[[FOR_RES]] : tensor<128x128xf32>, tensor<128x128xf32>, tensor<128x128xf32>
+  // CHECK-NEXT: %[[R1:.*]] = bufferization.to_tensor %[[M1]] : memref<128x128xf32>
+  // CHECK-NEXT: return %[[R0]], %[[R1]], %[[FOR_RES]] : tensor<128x128xf32>, tensor<128x128xf32>, tensor<128x128xf32>
   return %0#0, %0#1, %0#2 : tensor<128x128xf32>, tensor<128x128xf32>, tensor<128x128xf32>
 }
 
 // -----
 
 // CHECK-LABEL: fold_away_iter_with_no_use_and_yielded_input
 //  CHECK-SAME:   %[[A0:[0-9a-z]*]]: i32
 func.func @fold_away_iter_with_no_use_and_yielded_input(%arg0 : i32,

mlir/test/Dialect/SparseTensor/sparse_vector_chain.mlir

 // CHECK:               %[[VAL_82:.*]] = arith.addf %[[VAL_77]], %[[VAL_81]] : vector<8xf64>
 // CHECK:               %[[VAL_83:.*]] = arith.select %[[VAL_80]], %[[VAL_82]], %[[VAL_77]] : vector<8xi1>, vector<8xf64>
 // CHECK:               scf.yield %[[VAL_83]] : vector<8xf64>
 // CHECK:             }
 // CHECK:             %[[VAL_84:.*]] = vector.reduction <add>, %[[VAL_85:.*]] : vector<8xf64> into f64
 // CHECK:             scf.yield %[[VAL_84]] : f64
 // CHECK:           }
 // CHECK:           memref.store %[[VAL_86:.*]], %[[VAL_15]][] : memref<f64>
-// CHECK:           return %[[VAL_0]] : tensor<f64>
+// CHECK:           %[[VAL_87:.*]] = bufferization.to_tensor %[[VAL_15]] : memref<f64>
+// CHECK:           return %[[VAL_87]] : tensor<f64>
 // CHECK:         }
 func.func @sparse_matrix_sum(%argx: tensor<f64> {linalg.inplaceable = true},
                          %arga: tensor<64x32xf64, #SparseMatrix>,
                          %argb: tensor<64x32xf64, #SparseMatrix>) -> tensor<f64> {
   %0 = linalg.generic #trait
      ins(%arga, %argb: tensor<64x32xf64, #SparseMatrix>,
                        tensor<64x32xf64, #SparseMatrix>)
       outs(%argx: tensor<f64>) {
       ^bb(%a: f64, %b: f64, %x: f64):
         %m = arith.addf %a, %b : f64
         %t = arith.addf %x, %m : f64
         linalg.yield %t : f64
   } -> tensor<f64>
   return %0 : tensor<f64>
 }