[mlir][bufferization] Add read_only attribute to ToMemrefOp

[springerm]

Summary

This unit attribute indicates to the bufferization that the resulting buffer will not be written to by another op.

Depends On: D154966

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

let description = [{
     // Result type is memref<4x?xf32, #layout, 0>
     %m = bufferization.to_memref %t : memref<4x?xf32, #layout, 0>
     ```
 
     This operation is a specialized variant of the built-in
     `unrealized_conversion_cast` and is used to make sure that the IR stays
     valid at any point during the bufferization.
 
-    IR that contains `to_memref` ops cannot be bufferized with One-Shot
-    Bufferize.
+    The `read_only` attribute can optionally be set, indicating to the
+    bufferization that the buffer returned by this op (or an alias created from
+    the returned buffer) will not be written to.
   }];
 
-  let arguments = (ins AnyTensor:$tensor);
+  let arguments = (ins AnyTensor:$tensor, UnitAttr:$read_only);
   let results = (outs AnyRankedOrUnrankedMemRef:$memref);
 
   let extraClassDeclaration = [{
     //===------------------------------------------------------------------===//
     // BufferizableOpInterface implementation
     //===------------------------------------------------------------------===//
 
     // Note: ToMemrefOp / ToTensorOp are temporary ops that are inserted at the
     // bufferization boundary. When One-Shot bufferization is complete, there
     // should be no such ops left over. If `allowUnknownOps` (or after running a
     // partial bufferization pass), such ops may be part of the resulting IR,
     // but such IR may no longer be analyzable by One-Shot analysis.
 
     bool bufferizesToMemoryRead(OpOperand &opOperand,
                                 const AnalysisState &state) const {
       // It is unknown whether the resulting memref will be read or not.
       return true;
     }
 
     bool bufferizesToMemoryWrite(OpOperand &opOperand,
-                                 const AnalysisState &state) const {
-      // It is unknown whether the resulting MemRef will be written or not.
-      return true;
+                                 const AnalysisState &state) {
+      return !getReadOnly();
     }
 
     AliasingOpResultList getAliasingOpResults(
         OpOperand &opOperand, const AnalysisState &state) const {
       return {};
     }
 
     LogicalResult bufferize(RewriterBase &rewriter,
                             const BufferizationOptions &options);
   }];
 
-  let assemblyFormat = "$tensor attr-dict `:` type($memref)";
+  let assemblyFormat = [{
+    $tensor (`read_only` $read_only^)? attr-dict `:` type($memref)
+  }];
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 #endif // BUFFERIZATION_OPS

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

 func.func @read_of_alloc_tensor_is_not_a_conflict(%f: f32, %idx: index) -> f32 {
   %0 = bufferization.alloc_tensor() : tensor<10xf32>
   // This can be in-place because the read below does reads undefined data.
   // CHECK: tensor.insert {{.*}} {__inplace_operands_attr__ = ["none", "true", "none"]}
   %1 = tensor.insert %f into %0[%idx] : tensor<10xf32>
   %2 = tensor.extract %0[%idx] : tensor<10xf32>
   return %2 : f32
 }
+
+// -----
+
+// CHECK-LABEL: func @to_memref_not_read_only(
+func.func @to_memref_not_read_only(%idx : index, %f: f32) -> f32 {
+  %t = tensor.generate {
+  ^bb0(%i : index):
+    tensor.yield %f : f32
+  } : tensor<5xf32>
+  // Some op may write into the result of to_memref later.
+  // CHECK: bufferization.to_memref
+  // CHECK-SAME: {__inplace_operands_attr__ = ["false"]}
+  %m = bufferization.to_memref %t : memref<5xf32>
+  %2 = tensor.extract %t[%idx] : tensor<5xf32>
+  return %2 : f32
+}
+
+// -----
+
+// CHECK-LABEL: func @to_memref_read_only(
+func.func @to_memref_read_only(%idx : index, %f: f32) -> f32 {
+  %t = tensor.generate {
+  ^bb0(%i : index):
+    tensor.yield %f : f32
+  } : tensor<5xf32>
+  // Some op may write into the result of to_memref later.
+  // CHECK: bufferization.to_memref
+  // CHECK-SAME: {__inplace_operands_attr__ = ["true"]}
+  %m = bufferization.to_memref %t {read_only} : memref<5xf32>
+  %2 = tensor.extract %t[%idx] : tensor<5xf32>
+  return %2 : f32
+}