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 @@ -413,12 +413,6 @@ 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 {}; diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp --- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp @@ -938,13 +938,6 @@ 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(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. diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir --- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir @@ -231,23 +231,6 @@ // ----- -func.func @to_memref_op_unsupported( - %t1: tensor {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 - - // Read from both. - %cst = arith.constant 0.0 : f32 - %r1 = vector.transfer_read %t1[%idx3], %cst : tensor, vector<5xf32> - %r2 = vector.transfer_read %0[%idx3], %cst : memref, vector<5xf32> - - return %r1, %r2 : vector<5xf32>, vector<5xf32> -} - -// ----- - func.func @to_tensor_op_unsupported(%m: memref, %idx: index) -> (f32) { // expected-error @+1 {{to_tensor ops without `restrict` are not supported by One-Shot Analysis}} %0 = bufferization.to_tensor %m : memref diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir --- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir @@ -636,3 +636,26 @@ llvm.call @llvm_func() : () -> () return } + +// ----- + +// CHECK-LABEL: func @to_memref_op_unsupported( +// CHECK-SAME: %[[arg0:.*]]: memref {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 + // CHECK: "test.foo"(%[[alloc]]) + "test.foo"(%0) : (memref) -> () + + // CHECK: vector.transfer_read %[[arg0]] + %cst = arith.constant 0.0 : f32 + %r1 = vector.transfer_read %t1[%idx3], %cst : tensor, vector<5xf32> + + return %r1 : vector<5xf32> +}