[mlir] Add support for tensor.extract to comprehensive bufferization

[ThomasRaoux]

Summary

Depends on https://reviews.llvm.org/D105857

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

static LogicalResult bufferize(OpBuilder &b, linalg::YieldOp yieldOp,
   // No tensors -> success.
   if (!llvm::any_of(yieldOp.getOperandTypes(), isaTensor))
     return success();
   // linalg::YieldOp nested under TiledLoop must just canonicalize.
   if (yieldOp->getParentOfType<TiledLoopOp>())
     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<memref::LoadOp>(loc, srcMemref, extractOp.indices());
+  extractOp.replaceAllUsesWith(l);
+  return success();
+}
 //===----------------------------------------------------------------------===//
 // Bufferization analyses.
 //===----------------------------------------------------------------------===//
 
 ///
 /// Rationale for bufferizing `%1 = tensor.extract_slice %0[...]` inplace.
 /// ===========================================================
 ///
WalkResult result =
       .Case<memref::BufferCastOp,
             memref::TensorLoadOp>([&](auto) { return success(); })
       .Case<tensor::CastOp,
             tensor::DimOp,
             ExtractSliceOp,
             scf::ForOp,
             InitTensorOp,
             InsertSliceOp,
+            tensor::ExtractOp,
             LinalgOp,
             ReturnOp,
             TiledLoopOp,
             VectorTransferOpInterface,
             linalg::YieldOp,
             scf::YieldOp>([&](auto op) {
         LDBG("Begin bufferize:\n" << op << '\n');
         return bufferize(b, op, bvm, aliasInfo);

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

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