[mlir][Linalg] Fail comprehensive bufferization if a memref is returned.

[nicolasvasilache]

Add a flag to allow returning a memref for testing purposes.

[nicolasvasilache]

Rebase

mlir/include/mlir/Dialect/Linalg/Passes.td

let description = [{
     FuncOp within a module. It provides analyses and bufferization across
     function boundaries. Within a function boundary, the analysis is performed
     on SSA use-def chains starting from function operands that are annotated
     with the 'inplaceable' attribute.
   }];
   let options = [
     Option<"testAnalysisOnly", "test-analysis-only", "bool",
             /*default=*/"false",
-           "Only runs inplaceability analysis (for testing purposes only)">
+           "Only runs inplaceability analysis (for testing purposes only)">,
+    Option<"allowReturnMemref", "allow-return-memref", "bool",
+            /*default=*/"false",
+           "Allows the return of memrefs (for testing purposes only)">
   ];
   let constructor = "mlir::createLinalgComprehensiveModuleBufferizePass()";
 }
 
 def LinalgFoldUnitExtentDims : FunctionPass<"linalg-fold-unit-extent-dims"> {
   let summary = "Remove unit-extent dimension in Linalg ops on tensors";
   let constructor = "mlir::createLinalgFoldUnitExtentDimsPass()";
   let options = [

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp

void LinalgComprehensiveModuleBufferize::runOnOperation() {
   for (FuncOp funcOp : orderedFuncOps) {
     // Note: It would be good to apply cleanups here but we cannot as aliasInfo
     // would be invalidated.
     if (failed(bufferizeFuncOpBoundary(funcOp, aliasInfo,
                                        bufferizedFunctionTypes))) {
       signalPassFailure();
       return;
     }
+    if (!allowReturnMemref &&
+        llvm::any_of(funcOp.getType().getResults(), [](Type t) {
+          return t.isa<MemRefType, UnrankedMemRefType>();
+        })) {
+      funcOp->emitError("memref return type is unsupported");
+      signalPassFailure();
+      return;
+    }
   }
 
   // Perform a post-processing pass of layout modification at function boundary
   // according to the kBufferLayoutAttrName.
   layoutPostProcessing(moduleOp);
 
   // Post-pass cleanup of inplaceable and buffer_layout attributes.
   moduleOp.walk(

mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir

 // -----
 
 func @unknown_op(%A : tensor<4xf32>) -> tensor<4xf32>
 {
   // expected-error @+1 {{unsupported op with tensors}}
   %r = "marklar"(%A) : (tensor<4xf32>) -> (tensor<4xf32>)
   return %r: tensor<4xf32>
 }
+
+// -----
+
+// expected-error @+1 {{memref return type is unsupported}}
+func @mini_test_case1() -> tensor<10x20xf32> {
+  %f0 = constant 0.0 : f32
+  %t = linalg.init_tensor [10, 20] : tensor<10x20xf32>
+  %r = linalg.fill(%f0, %t) : f32, tensor<10x20xf32> -> tensor<10x20xf32>
+  return %r : tensor<10x20xf32>
+}

mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir

-// RUN: mlir-opt %s -linalg-comprehensive-module-bufferize -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -linalg-comprehensive-module-bufferize=allow-return-memref -split-input-file | FileCheck %s
 
 // CHECK-LABEL: func @transfer_read(%{{.*}}: memref<?xf32, #map>) -> vector<4xf32> {
 func @transfer_read(%A : tensor<?xf32>) -> (vector<4xf32>) {
   %c0 = constant 0 : index
   %f0 = constant 0.0 : f32
 
 //       CHECK: %[[RES:.*]] = vector.transfer_read {{.*}} : memref<?xf32, #{{.*}}>, vector<4xf32>
   %0 = vector.transfer_read %A[%c0], %f0 : tensor<?xf32>, vector<4xf32>