[mlir][bufferization] Allow to_memref ops in One-Shot Analysis

[springerm]

Summary

bufferization.to_memref ops are allowed in One-Shot Bufferize, but they are treated conservatively: in the absence of a memref analysis, we have to assume that the result buffer is read and written.

Note: to_memref cannot introduce any future aliases that would have to be considered during One-Shot Bufferize, because only to_tensor ops with the restrict attribute are supported. Such tensors are guaranteed to not alias with any other buffer after bufferization.

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td

let extraClassDeclaration = [{
     }
 
     bool bufferizesToMemoryWrite(OpOperand &opOperand,
                                  const AnalysisState &state) const {
       // It is unknown whether the resulting MemRef will be written or not.
       return true;
     }
 
-    bool mustBufferizeInPlace(OpOperand &opOperand,
-                              const AnalysisState &state) const {
-      // ToMemrefOps always bufferize inplace.
-      return true;
-    }
-
     AliasingOpResultList getAliasingOpResults(
         OpOperand &opOperand, const AnalysisState &state) const {
       return {};
     }
 
     LogicalResult bufferize(RewriterBase &rewriter,
                             const BufferizationOptions &options);
   }];
 
   let assemblyFormat = "$tensor attr-dict `:` type($memref)";
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 #endif // BUFFERIZATION_OPS

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

static LogicalResult checkAliasInfoConsistency(Operation *op,
                                                OneShotAnalysisState &state) {
   const BufferizationOptions &options = state.getOptions();
 
   WalkResult walkResult = op->walk([&](BufferizableOpInterface op) {
     // Skip ops that are not in the filter.
     if (!options.isOpAllowed(op.getOperation()))
       return WalkResult::advance();
 
-    // Input IR may not contain any ToMemrefOps. These are not supported because
-    // the analysis cannot follow the data flow through memrefs.
-    if (isa<ToMemrefOp>(op.getOperation())) {
-      op->emitError("to_memref ops are not supported by One-Shot Analysis");
-      return WalkResult::interrupt();
-    }
-
     // Input IR may not contain any ToTensorOps without the "restrict"
     // attribute. Such tensors may alias any other tensor, which is currently
     // not handled in the analysis.
     if (auto toTensorOp = dyn_cast<ToTensorOp>(op.getOperation())) {
       if (!toTensorOp.getRestrict()) {
         op->emitError("to_tensor ops without `restrict` are not supported by "
                       "One-Shot Analysis");
         return WalkResult::interrupt();

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir

func.func @main() -> tensor<4xi32> {
   }
 
   // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}}
   return %r: tensor<4xi32>
 }
 
 // -----
 
-func.func @to_memref_op_unsupported(
-    %t1: tensor<?xf32> {bufferization.writable = true}, %idx1: index,
-    %idx2: index, %idx3: index, %v1: vector<5xf32>) -> (vector<5xf32>, vector<5xf32>) {
-
-  // expected-error @+1 {{to_memref ops are not supported by One-Shot Analysis}}
-  %0 = bufferization.to_memref %t1 : memref<?xf32>
-
-  // Read from both.
-  %cst = arith.constant 0.0 : f32
-  %r1 = vector.transfer_read %t1[%idx3], %cst : tensor<?xf32>, vector<5xf32>
-  %r2 = vector.transfer_read %0[%idx3], %cst : memref<?xf32>, vector<5xf32>
-
-  return %r1, %r2 : vector<5xf32>, vector<5xf32>
-}
-
-// -----
-
 func.func @to_tensor_op_unsupported(%m: memref<?xf32>, %idx: index) -> (f32) {
   // expected-error @+1 {{to_tensor ops without `restrict` are not supported by One-Shot Analysis}}
   %0 = bufferization.to_tensor %m : memref<?xf32>
 
   %1 = tensor.extract %0[%idx] : tensor<?xf32>
   return %1 : f32
 }
 

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir

 // Test if other callables are left intact and don't cause trouble.
 
 llvm.func @llvm_func()
 
 func.func @call_llvm_func() {
   llvm.call @llvm_func() : () -> ()
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @to_memref_op_unsupported(
+//  CHECK-SAME:     %[[arg0:.*]]: memref<?xf32,
+func.func @to_memref_op_unsupported(
+    %t1: tensor<?xf32> {bufferization.writable = true}, %idx1: index,
+    %idx2: index, %idx3: index, %v1: vector<5xf32>) -> (vector<5xf32>) {
+
+  // Insert a copy because we cannot analyze what happens with the result of a
+  // to_memref op.
+  // CHECK: %[[alloc:.*]] = memref.alloc
+  // CHECK: memref.copy %[[arg0]], %[[alloc]]
+  %0 = bufferization.to_memref %t1 : memref<?xf32>
+  // CHECK: "test.foo"(%[[alloc]])
+  "test.foo"(%0) : (memref<?xf32>) -> ()
+
+  // CHECK: vector.transfer_read %[[arg0]]
+  %cst = arith.constant 0.0 : f32
+  %r1 = vector.transfer_read %t1[%idx3], %cst : tensor<?xf32>, vector<5xf32>
+
+  return %r1 : vector<5xf32>
+}