diff --git a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp @@ -2379,6 +2379,8 @@ static BlockArgument getEquivalentEnclosingFuncBBArg(Value v, const BufferizationAliasInfo &aliasInfo) { + if (!v.getType().isa()) + return nullptr; Operation *op = v.getParentBlock()->getParentOp(); FuncOp funcOp = dyn_cast(op); if (!funcOp) @@ -2455,6 +2457,12 @@ // 1. For each FuncOp result, keep track of which inplace argument it reuses. SmallVector returnValues; for (OpOperand &returnOperand : returnOp->getOpOperands()) { + // If not a renturn tensor type just forward it. + if (!returnOperand.get().getType().isa()) { + 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)) diff --git a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir --- a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir +++ b/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) -> vector<4xf32> { +func @transfer_read(%A : tensor) -> (vector<4xf32>) { + %c0 = constant 0 : index + %f0 = constant 0.0 : f32 + +// CHECK: %[[RES:.*]] = vector.transfer_read {{.*}} : memref, vector<4xf32> + %0 = vector.transfer_read %A[%c0], %f0 : tensor, 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(