Index: mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp =================================================================== --- mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp +++ mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp @@ -2113,6 +2113,22 @@ return success(); llvm_unreachable("unexpected yieldOp"); } + +/// Bufferization for tensor::ExtractOp just translate to memref.load, it only +/// reads the tensor. +static LogicalResult bufferize(OpBuilder &b, tensor::ExtractOp extractOp, + BlockAndValueMapping &bvm, + BufferizationAliasInfo &aliasInfo) { + // Take a guard before anything else. + OpBuilder::InsertionGuard g(b); + b.setInsertionPoint(extractOp); + + Location loc = extractOp.getLoc(); + Value srcMemref = lookup(bvm, extractOp.tensor()); + Value l = b.create(loc, srcMemref, extractOp.indices()); + extractOp.replaceAllUsesWith(l); + return success(); +} //===----------------------------------------------------------------------===// // Bufferization analyses. //===----------------------------------------------------------------------===// @@ -2308,6 +2324,7 @@ scf::ForOp, InitTensorOp, InsertSliceOp, + tensor::ExtractOp, LinalgOp, ReturnOp, TiledLoopOp, Index: mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir =================================================================== --- mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir +++ mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir @@ -20,6 +20,19 @@ // ----- +// CHECK-LABEL: func @tensor_extract(%{{.*}}: memref) -> f32 { +func @tensor_extract(%A : tensor) -> (f32) { + %c0 = constant 0 : index + +// CHECK: %[[RES:.*]] = memref.load {{.*}} : memref + %0 = tensor.extract %A[%c0] : tensor + +// CHECK: return %[[RES]] : f32 + return %0 : f32 +} + +// ----- + // CHECK-DAG: #[[$map_1d_dyn:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)> /// No linalg.inplaceable flag, must allocate.