diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td --- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td +++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td @@ -426,11 +426,12 @@ `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 @@ } 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 @@ 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 --- 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 @@ %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 +}