[mlir][Linalg] Better support for bufferizing non-tensor results.

Clean up corner cases related to elemental tensor / buffer type return values that would previously fail.

Reviewed By: ThomasRaoux

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

Author: nicolasvasilache

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

@@ -2379,6 +2379,8 @@ static Operation *getEquivalentAlloc(Value value,
 static BlockArgument
 getEquivalentEnclosingFuncBBArg(Value v,
                                 const BufferizationAliasInfo &aliasInfo) {
+  if (!v.getType().isa<RankedTensorType>())
+    return nullptr;
   Operation *op = v.getParentBlock()->getParentOp();
   FuncOp funcOp = dyn_cast<FuncOp>(op);
   if (!funcOp)
@@ -2455,6 +2457,12 @@ static LogicalResult bufferizeFuncOpBoundary(
   // 1. For each FuncOp result, keep track of which inplace argument it reuses.
   SmallVector<Value> returnValues;
   for (OpOperand &returnOperand : returnOp->getOpOperands()) {
+    // If not a renturn tensor type just forward it.
+    if (!returnOperand.get().getType().isa<RankedTensorType>()) {
+      returnValues.push_back(returnOperand.get());
+      continue;
+    }
+
     // If return operand is equivalent to some bbArg, no need to return it.
     Value returnVal = returnOperand.get();
     if (getEquivalentEnclosingFuncBBArg(returnVal, aliasInfo))

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

@@ -1,5 +1,19 @@
 // RUN: mlir-opt %s -linalg-comprehensive-module-bufferize -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>
+
+//       CHECK: return %[[RES]] : vector<4xf32>
+  return %0 : vector<4xf32>
+}
+
+// -----
+
 // CHECK-DAG: #[[$map_1d_dyn:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
 
 // CHECK-LABEL: func @fill_inplace(