[mlir][bufferize] Insert memref.cast ops during finalizing pass

[springerm]

Summary

The pass can currently not handle to_memref(to_tensor(x)) folding where a cast is necessary. This is required with the new unified bufferization. There is already a canonicalization pattern that handles such foldings and it should be used during this pass.

Depends On D118243

[nicolasvasilache]

This is scary on many levels, we don't want to propagate this.
Can we turn this into an assertion and be very loud if something does not comply?

[nicolasvasilache]

Ok this thing is actually never used in the preferred case, it can only trigger if one uses the old bufferization mixed in which should be deprecated in the future.

[springerm]

add test case

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

 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 
 //===----------------------------------------------------------------------===//
 // Bufferization Dialect
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Bufferization/IR/BufferizationOpsDialect.h.inc"
 
+namespace mlir {
+class RewritePatternSet;
+class MLIRContext;
+
+namespace bufferization {
+/// Populate patterns for folding to_memref and to_tensor ops.
+/// Note: to_memref(to_tensor(x)) without type changes are handled by a folder.
+void populateBufferizationOpFoldingPatterns(RewritePatternSet &patterns,
+                                            MLIRContext *context);
+} // namespace bufferization
+} // namespace mlir
+
 //===----------------------------------------------------------------------===//
 // Bufferization Dialect Operations
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Bufferization/IR/BufferizationOps.h.inc"
 
 #endif // MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZATION_H_

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

if (memrefToTensorType && toMemrefType &&
       !isGuaranteedCastCompatible(memrefToTensorType, toMemrefType)) {
     MemRefType resultType = toMemrefType;
     auto loc = toMemref.getLoc();
     SmallVector<Value, 4> dynamicOperands;
     for (int i = 0; i < resultType.getRank(); ++i) {
       if (resultType.getShape()[i] != ShapedType::kDynamicSize)
         continue;
       auto index = rewriter.createOrFold<arith::ConstantIndexOp>(loc, i);
-      Value size = rewriter.create<tensor::DimOp>(loc, memrefToTensor, index);
+      Value size =
+          rewriter.create<memref::DimOp>(loc, memrefToTensor.memref(), index);
       dynamicOperands.push_back(size);
     }
     // TODO: Use alloc/memcpy callback from BufferizationOptions if called via
     // BufferizableOpInterface impl of ToMemrefOp.
     auto copy =
         rewriter.create<memref::AllocOp>(loc, resultType, dynamicOperands);
     rewriter.create<memref::CopyOp>(loc, memrefToTensor.memref(), copy);
     rewriter.replaceOp(toMemref, {copy});

 } // namespace
 
 void ToMemrefOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
   results.add<DimOfCastOp, LoadOfToMemref, ToMemrefOfCast, TensorLoadToMemref>(
       context);
 }
 
+void bufferization::populateBufferizationOpFoldingPatterns(
+    RewritePatternSet &patterns, MLIRContext *context) {
+  patterns.add<TensorLoadToMemref>(context);
+}
+
 LogicalResult ToMemrefOp::bufferize(RewriterBase &rewriter,
                                     const BufferizationState &state) {
   // Fold to_memref(to_tensor(x)) to x. Insert a cast if necessary.
   return foldToMemrefToTensorPair(rewriter, *this);
 }
 
 Optional<Operation *> CloneOp::buildDealloc(OpBuilder &builder, Value alloc) {
   return builder.create<memref::DeallocOp>(alloc.getLoc(), alloc)

mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp

public:
   }
 };
 } // namespace
 
 void mlir::bufferization::populateEliminateBufferizeMaterializationsPatterns(
     BufferizeTypeConverter &typeConverter, RewritePatternSet &patterns) {
   patterns.add<BufferizeToTensorOp, BufferizeToMemrefOp>(typeConverter,
                                                          patterns.getContext());
+  populateBufferizationOpFoldingPatterns(patterns, patterns.getContext());
 }
 
 namespace {
 struct FinalizingBufferizePass
     : public FinalizingBufferizeBase<FinalizingBufferizePass> {
   using FinalizingBufferizeBase<
       FinalizingBufferizePass>::FinalizingBufferizeBase;
 

mlir/test/Dialect/Bufferization/Transforms/finalizing-bufferize.mlir

 // RUN: mlir-opt %s -finalizing-bufferize -split-input-file -verify-diagnostics | FileCheck %s
 
-// CHECK-LABEL:   func @eliminate_materializations(
-// CHECK-SAME:                                     %[[ARG:.*]]: memref<f32>) -> memref<f32> {
-// CHECK:           return %[[ARG]] : memref<f32>
+// CHECK-LABEL: func @eliminate_materializations(
+//  CHECK-SAME:     %[[ARG:.*]]: memref<f32>) -> memref<f32> {
 func @eliminate_materializations(%arg0: memref<f32>) -> memref<f32> {
   %0 = bufferization.to_tensor %arg0 : memref<f32>
   %1 = bufferization.to_memref %0 : memref<f32>
+  // CHECK: return %[[ARG]] : memref<f32>
   return %1 : memref<f32>
 }
 
 // -----
 
 func @unable_to_convert_lone_buffer_cast() -> memref<f32> {
   // expected-error @+1 {{failed to legalize operation 'test.source'}}
   %0 = "test.source"() : () -> tensor<f32>
   %1 = bufferization.to_memref %0 : memref<f32>
   return %1 : memref<f32>
 }
 
 // -----
 
 func @unable_to_convert_lone_tensor_load(%arg0: memref<f32>) {
   %0 = bufferization.to_tensor %arg0 : memref<f32>
   // expected-error @+1 {{failed to legalize operation 'test.sink'}}
   "test.sink"(%0) : (tensor<f32>) -> ()
   return
 }
+
+// -----
+
+// CHECK: #[[$MAP1:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+#map1 = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+
+// CHECK-LABEL: func @insert_memref_cast(
+//  CHECK-SAME:     %[[arg0:.*]]: memref<?xf32>
+func @insert_memref_cast(%arg0: memref<?xf32>) -> memref<?xf32, #map1> {
+  %0 = bufferization.to_tensor %arg0 : memref<?xf32>
+  %1 = bufferization.to_memref %0 : memref<?xf32, #map1>
+  // CHECK: %[[r:.*]] = memref.cast %[[arg0]] : memref<?xf32> to memref<?xf32, #[[$MAP1]]>
+  // CHECK: return %[[r]]
+  return %1 : memref<?xf32, #map1>
+}
+
+// -----
+
+// CHECK: #[[$MAP2:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+#map2 = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+
+// CHECK-LABEL: func @insert_buffer_copy(
+//  CHECK-SAME:     %[[arg0:.*]]: memref<?xf32, #[[$MAP2]]>
+func @insert_buffer_copy(%arg0: memref<?xf32, #map2>) -> memref<?xf32> {
+  // CHECK: %[[c0:.*]] = arith.constant 0 : index
+  // CHECK: %[[dim0:.*]] = memref.dim %[[arg0]], %[[c0]]
+  // CHECK: %[[alloc:.*]] = memref.alloc(%[[dim0]]) : memref<?xf32>
+  // CHECK: memref.copy %[[arg0]], %[[alloc]]
+  %0 = bufferization.to_tensor %arg0 : memref<?xf32, #map2>
+  %1 = bufferization.to_memref %0 : memref<?xf32>
+
+  // CHECK: return %[[alloc]]
+  return %1 : memref<?xf32>
+}