[Mlir-commits] [mlir] 8e72fbd - [mlir][bufferization] Add read_only attribute to ToMemrefOp

[Matthias Springer]

Author: Matthias Springer
Date: 2023-07-11T16:37:17+02:00
New Revision: 8e72fbd616ccc98e489aee634b07747ae46fd8b6

URL: https://github.com/llvm/llvm-project/commit/8e72fbd616ccc98e489aee634b07747ae46fd8b6
DIFF: https://github.com/llvm/llvm-project/commit/8e72fbd616ccc98e489aee634b07747ae46fd8b6.diff

LOG: [mlir][bufferization] Add read_only attribute to ToMemrefOp

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

Differential Revision: https://reviews.llvm.org/D154967

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
    mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
index a1e2f8114269fb..ef10b9a65ab081 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
@@ -426,11 +426,12 @@ def Bufferization_ToMemrefOp : Bufferization_Op<"to_memref", [
     `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 = [{
@@ -451,9 +452,8 @@ def Bufferization_ToMemrefOp : Bufferization_Op<"to_memref", [
     }
 
     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(
@@ -465,7 +465,9 @@ def Bufferization_ToMemrefOp : Bufferization_Op<"to_memref", [
                             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;

diff  --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir
index df174d01ca07a0..0938f3d2dd17d4 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir
@@ -66,3 +66,35 @@ func.func @read_of_alloc_tensor_is_not_a_conflict(%f: f32, %idx: index) -> f32 {
   %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
+}