[mlir][bufferize] Improve to_tensor/to_memref folding

[springerm]

Summary

Fold all to_tensor/to_memref pairs during canonicalization. This is needed for some sparse compiler test cases, but will also lead to cleaner IR post-bufferization in general.

mlir/include/mlir/Dialect/Bufferization/IR/Bufferization.h

 /// This function returns `failure()` in case of unsupported casts. E.g., casts
 /// with differing element types or memory spaces.
 FailureOr<Value> castOrReallocMemRefValue(OpBuilder &b, Value value,
                                           MemRefType type);
 
 /// Try to fold to_memref(to_tensor(x)). If x's type and the result type of the
 /// to_memref op are different, a memref.cast is needed.
 LogicalResult foldToMemrefToTensorPair(RewriterBase &rewriter,
-                                       ToMemrefOp toMemref,
-                                       bool allowSameType = true);
+                                       ToMemrefOp toMemref);
 
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZATION_H_

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

mlir::bufferization::castOrReallocMemRefValue(OpBuilder &b, Value value,
   // BufferizableOpInterface impl of ToMemrefOp.
   Value copy = b.create<memref::AllocOp>(loc, destType, dynamicOperands);
   b.create<memref::CopyOp>(loc, value, copy);
   return copy;
 }
 
 /// Try to fold to_memref(to_tensor(x)). If x's type and the result type of the
 /// to_memref op are different, a memref.cast is needed.
-LogicalResult mlir::bufferization::foldToMemrefToTensorPair(
-    RewriterBase &rewriter, ToMemrefOp toMemref, bool allowSameType) {
+LogicalResult
+mlir::bufferization::foldToMemrefToTensorPair(RewriterBase &rewriter,
+                                              ToMemrefOp toMemref) {
   auto memrefToTensor = toMemref.getTensor().getDefiningOp<ToTensorOp>();
   if (!memrefToTensor)
     return failure();
 
   Type srcType = memrefToTensor.getMemref().getType();
   Type destType = toMemref.getType();
 
   // Directly rewrite if the type did not change.
   if (srcType == destType) {
-    // Function can be configured to only handle cases where a cast is needed.
-    if (!allowSameType)
-      return failure();
     rewriter.replaceOp(toMemref, memrefToTensor.getMemref());
     return success();
   }
 
   auto rankedSrcType = srcType.dyn_cast<MemRefType>();
   auto rankedDestType = destType.dyn_cast<MemRefType>();
   auto unrankedSrcType = srcType.dyn_cast<UnrankedMemRefType>();
 
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.source().getDefiningOp<ToTensorOp>();
     if (!memrefToTensorOp)
       return failure();
 
     rewriter.replaceOpWithNewOp<memref::DimOp>(
         dimOp, memrefToTensorOp.getMemref(), dimOp.index());
     return success();
   }
 };
-
 } // namespace
 
 void ToTensorOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
-  results.add<DimOfToTensorFolder>(context);
+  results.add<DimOfToTensorFolder, ToTensorToMemrefFolding>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // ToMemrefOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult ToMemrefOp::fold(ArrayRef<Attribute>) {
   if (auto memrefToTensor = getTensor().getDefiningOp<ToTensorOp>())
LogicalResult matchAndRewrite(ToMemrefOp toMemref,
     Value memref = rewriter.create<ToMemrefOp>(toMemref.getLoc(), memrefType,
                                                tensorCastOperand.getOperand());
     rewriter.replaceOpWithNewOp<memref::CastOp>(toMemref, toMemref.getType(),
                                                 memref);
     return success();
   }
 };
 
-/// Canonicalize bufferization.to_tensor + bufferization.to_memref to
-/// memref.cast when type mismatches prevent `ToMemrefOp::fold` to kick in.
-struct TensorLoadToMemref : public OpRewritePattern<ToMemrefOp> {
+/// Canonicalize bufferization.to_tensor + bufferization.to_memref. Insert a
+/// cast if necessary.
+struct ToMemrefToTensorFolding : public OpRewritePattern<ToMemrefOp> {
   using OpRewritePattern<ToMemrefOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ToMemrefOp toMemref,
                                 PatternRewriter &rewriter) const final {
-    // Only handle cases where a cast is needed. The other case is handled by
-    // the folder.
-    return foldToMemrefToTensorPair(rewriter, toMemref,
-                                    /*allowSameType=*/false);
+    return foldToMemrefToTensorPair(rewriter, toMemref);
   }
 };
 
 /// Fold a load on a to_memref operation into an tensor.extract on the
 /// corresponding tensor.
 struct LoadOfToMemref : public OpRewritePattern<memref::LoadOp> {
   using OpRewritePattern<memref::LoadOp>::OpRewritePattern;
 
LogicalResult matchAndRewrite(memref::DimOp dimOp,
     return success();
   }
 };
 
 } // namespace
 
 void ToMemrefOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
-  results.add<DimOfCastOp, LoadOfToMemref, ToMemrefOfCast, TensorLoadToMemref>(
-      context);
+  results.add<DimOfCastOp, LoadOfToMemref, ToMemrefOfCast,
+              ToMemrefToTensorFolding>(context);
 }
 
 LogicalResult ToMemrefOp::bufferize(RewriterBase &rewriter,
                                     const BufferizationOptions &options) {
   // Fold to_memref(to_tensor(x)) to x. Insert a cast if necessary.
   (void)foldToMemrefToTensorPair(rewriter, *this);
   // Note: The return value of `bufferize` indicates whether there was an error
   // or not. (And not whether the pattern matched or not.)

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: %[[R1:.*]] = bufferization.to_tensor %[[M1]] : memref<128x128xf32>
-  // CHECK-NEXT: return %[[R0]], %[[R1]], %[[FOR_RES]] : tensor<128x128xf32>, tensor<128x128xf32>, tensor<128x128xf32>
+  // CHECK-NEXT: return %[[R0]], %[[T1]], %[[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:           %[[VAL_87:.*]] = bufferization.to_tensor %[[VAL_15]] : memref<f64>
-// CHECK:           return %[[VAL_87]] : tensor<f64>
+// CHECK:           return %[[VAL_0]] : 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>
 }