[mlir][Linalg] Add comprehensive bufferization support for subtensor (5/n)

[nicolasvasilache]

Summary

This revision refactors and simplifies the pattern detection logic: thanks to SSA value properties, we can actually look at all the uses of a given value and avoid having to pattern-match specific chains of operations.

A bufferization pattern for subtensor is added and specific inplaceability analysis is implemented for the simple case of subtensor. More advanced use cases will follow.

[herhut]

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

771

nit: if

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

232

Should there be a test for this?

[nicolasvasilache]

Update.

[nicolasvasilache]

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

232

Added more test for the simple cases, a bunch of other steps are necessary to get everything in the right place, these will be coming in followup CLs.

[ftynse]

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

3

Why do we need tee here, looks like debug leftover.

It may not be available on Windows.

[nicolasvasilache]

Address comment.

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

 /// Determine which results may be reused inplace by the bufferization
 /// patterns of `bufferizeFuncOpInternals`.
 /// The inplace analysis uses this information along with interfering read
 /// analysis to determine which op results reuse the same buffer as some
 /// operand.
 OpResult getMatchingOpResult(OpOperand &opOperand) {
   OpResult res =
       llvm::TypeSwitch<Operation *, OpResult>(opOperand.getOwner())
-          .Case<LinalgOp, SubTensorInsertOp, VectorTransferOpInterface>(
+          // clang-format off
+          .Case<LinalgOp,
+                SubTensorInsertOp,
+                VectorTransferOpInterface>(
               [&](auto op) { return getMatchingOpResult(op, opOperand); })
+          .Case<SubTensorOp>([&](auto op) { return OpResult(); })
           .Default([&](Operation *op) { return OpResult(); });
+  // clang-format on
   return res;
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific attribute manipulation.
 //===----------------------------------------------------------------------===//
 
 /// Attribute marker to specify op results that can be bufferized inPlace.
for (OpOperand &operand : returnOp->getOpOperands()) {
     Value v = lookup(bvm, operand.get());
     if (!v)
       return failure();
     operand.set(b.create<memref::TensorLoadOp>(returnOp.getLoc(), v));
   }
   return success();
 }
 
+/// Bufferize SubTensorOp to subview with optional alloc + copy depending on
+/// whether or not it is marked inplaceable.
+/// Note that `getMatchingOpResult` on a SubTensorOp always returns null.
+/// As consequence a SubTensorOp always alloc + copy when taken in isolation.
+/// SubTensorOp can only be bufferized inplace when it appears with a companion
+/// SubTensorInsertOp.
+static LogicalResult bufferize(OpBuilder &b, SubTensorOp subTensorOp,
+                               BlockAndValueMapping &bvm) {
+  LLVM_DEBUG(DBGS() << "bufferize: " << *subTensorOp << "\n");
+
+  // Take a guard before anything else.
+  OpBuilder::InsertionGuard g(b);
+  b.setInsertionPoint(subTensorOp);
+
+  Location loc = subTensorOp.getLoc();
+  // Bail is source was not bufferized.
+  Value srcMemref = lookup(bvm, subTensorOp.source());
+  if (!srcMemref)
+    return failure();
+  auto srcMemrefType = srcMemref.getType().cast<MemRefType>();
+  auto dstTensorType = subTensorOp.result().getType().cast<RankedTensorType>();
+
+  // If not inplaceable, alloc.
+  Value alloc;
+  auto inPlace = getInPlace(subTensorOp->getResult(0));
+  if (inPlace != InPlaceSpec::True) {
+    alloc =
+        createNewAllocDeallocPairForShapedValue(b, loc, subTensorOp.result());
+    b.setInsertionPointAfter(alloc.getDefiningOp());
+  }
+
+  // Bufferize to subview.
+  auto subviewMemRefType =
+      memref::SubViewOp::inferRankReducedResultType(
+          dstTensorType.getRank(), srcMemrefType, subTensorOp.getMixedOffsets(),
+          subTensorOp.getMixedSizes(), subTensorOp.getMixedStrides())
+          .cast<MemRefType>();
+  Value subView = b.create<memref::SubViewOp>(
+      loc, subviewMemRefType, srcMemref, subTensorOp.getMixedOffsets(),
+      subTensorOp.getMixedSizes(), subTensorOp.getMixedStrides());
+
+  /// If not inplaceable, copy.
+  if (alloc) {
+    b.create<CopyOp>(subTensorOp.getLoc(), subView, alloc);
+    subView = alloc;
+  }
+
+  map(bvm, subTensorOp.result(), subView);
+  return success();
+}
 
 static LogicalResult bufferize(OpBuilder &b,
                                SubTensorInsertOp subTensorInsertOp,
                                BlockAndValueMapping &bvm) {
   LLVM_DEBUG(DBGS() << "bufferize: " << *subTensorInsertOp << "\n");
 
   // Take a guard before anything else.
   OpBuilder::InsertionGuard g(b);
static LogicalResult bufferizeFuncOpInternals(
   OpBuilder b(funcOp->getContext());
   /// Start by bufferizing `funcOp` arguments.
   if (failed(bufferize(b, funcOp, bvm)))
     return failure();
   WalkResult result = funcOp.walk<WalkOrder::PreOrder>([&](Operation *op) {
     LogicalResult status =
         llvm::TypeSwitch<Operation *, LogicalResult>(op)
             // Skip BufferCast and TensorLoad ops.
-            .Case<memref::BufferCastOp, memref::TensorLoadOp>(
+            // clang-format off
+            .Case<memref::BufferCastOp,
+                  memref::TensorLoadOp>(
                 [&](auto) { return success(); })
-            .Case<memref::DimOp, LinalgOp, ReturnOp, SubTensorInsertOp,
+            .Case<memref::DimOp,
+                  LinalgOp,
+                  ReturnOp,
+                  SubTensorOp,
+                  SubTensorInsertOp,
                   VectorTransferOpInterface>(
                 [&](auto op) { return bufferize(b, op, bvm); })
+            // clang-format on
             .Default([&](Operation *op) {
               auto isaTensor = [](Type t) { return t.isa<TensorType>(); };
               if (llvm::any_of(op->getOperandTypes(), isaTensor) ||
                   llvm::any_of(op->getResultTypes(), isaTensor))
                 return failure();
               return success();
             });
     if (failed(status)) {

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

 // RUN: mlir-opt %s -linalg-comprehensive-func-bufferize -split-input-file | FileCheck %s
 
 // CHECK-DAG: #[[$map_2d_dyn:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
 
 // CHECK-LABEL: func @fill_inplace(
 //  CHECK-SAME:   %[[A:[a-zA-Z0-9]*]]: tensor<?xf32> {linalg.inplaceable = true})
 func @fill_inplace(%A : tensor<?xf32> {linalg.inplaceable = true}) -> tensor<?xf32> {
   //     CHECK: %[[I:.*]] = memref.buffer_cast %[[A]] : memref<?xf32, #[[$map_2d_dyn]]>
 
   //     CHECK: %[[F0:.*]] = constant 0.000000e+00 : f32
   %f0 = constant 0.0 : f32
func @subtensor_insert_fun_not_inplace(%A : tensor<?xf32> {linalg.inplaceable = true}, %t : tensor<4xf32>)
   %r0 = subtensor_insert %t into %A[0][4][1] : tensor<4xf32> into tensor<?xf32>
 
   // TODO: WAW optimization where result is overwritten without being read.
   //      CHECK: linalg.fill(%[[BUFFER_CAST_A]]
   //      CHECK: memref.dealloc %[[ALLOC]] : memref<?xf32>
   %r1 = linalg.fill(%A, %f0) : tensor<?xf32>, f32 -> tensor<?xf32>
   return %r0, %r1: tensor<?xf32>, tensor<?xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @subtensor_fun
+func @subtensor_fun(%A : tensor<?xf32> {linalg.inplaceable = true})
+  ->  tensor<4xf32>
+{
+  //      CHECK: %[[BUFFER_CAST_A:.*]] = memref.buffer_cast {{.*}} : memref<?xf32
+
+  // In isolation, a subtensor will always alloc + copy.
+  // subtensor can only be bufferized inplace when it appears with a companion
+  // subtensor_insert
+  // CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<4xf32>
+  // CHECK: %[[SV:.*]] = memref.subview %[[BUFFER_CAST_A]][0] [4] [1]
+  // CHECK: linalg.copy(%[[SV]], %[[ALLOC]])
+  %r0 = subtensor %A[0][4][1] : tensor<?xf32> to tensor<4xf32>
+  return %r0: tensor<4xf32>
+}
+