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

[nicolasvasilache]

Summary

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

[ThomasRaoux]

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

3

nit: should be:
// CHECK-LABEL: func @transfer_read(%{{.*}}: memref<?xf32, #{{.*}}>) -> vector<4xf32> {

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

static Operation *getEquivalentAlloc(Value value,
   return res;
 }
 
 /// Return the first argument of the enclosing FuncOp that is equivalent to `v`.
 /// Return null if no such bbArg can be found.
 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)
     funcOp = op->getParentOfType<FuncOp>();
   assert(funcOp && "expected non-null FuncOp");
   for (BlockArgument bbArg : funcOp.getArguments()) {
     if (!bbArg.getType().isa<RankedTensorType>())
       continue;
static LogicalResult bufferizeFuncOpBoundary(
 
   // Support only single return-terminated block in the function.
   ReturnOp returnOp = getAssumedUniqueReturnOp(funcOp);
   assert(returnOp && "expected func with single return op");
 
   // 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))
       continue;
 
     // TODO: Need to hoist above function boundary.
     if (Operation *allocOp = getEquivalentAlloc(returnVal, aliasInfo)) {
       returnValues.push_back(allocOp->getResult(0));

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

 // 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(
 //  CHECK-SAME:   %[[A:[a-zA-Z0-9]*]]: memref<?xf32, #[[$map_1d_dyn]]>
 func @fill_inplace(%A : tensor<?xf32> {linalg.inplaceable = true}) -> tensor<?xf32> {
   //     CHECK: %[[F0:.*]] = constant 0.000000e+00 : f32
   %f0 = constant 0.0 : f32