[mlir][linalg] BufferizeToAllocationOp: Support vector.mask

[springerm]

Summary

This op needs special handling because the allocation for the masked op must be placed outside of the mask op.

Depends On: D153971

[dcaballe]

Thanks!

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

let description = [{
 
     ```
     %alloc = memref.alloc() : memref<10xf32>
     memref.tensor_store %dest, %alloc : memref<10xf32>
     memref.store %f, %alloc[%pos] : memref<10xf32>
     %0 = bufferization.to_tensor %alloc restrict writable : memref<10xf32>
     ```
 
-    Selected ops that bufferize to an allocation are also supported:
+    Selected ops that bufferize to an allocation (or need special handling) are
+    also supported:
     - `tensor.pad` is lowered to an allocation, followed by a `linalg.fill` and
       and a buffer copy (all on memrefs).
+    - `vector.mask` is bufferized together with its region. The allocation is
+      placed in front of the `vector.mask` op.
 
     An optional memory space attribute can be specified for the materialized
     buffer allocation.
 
     #### Return modes
 
     This operation consumes the `target` handle and produces the
     `allocated_buffer` handle. It always succeeds.

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 ///
 /// %alloc = memref.alloc
 /// linalg.fill ... outs(%alloc)
 /// %subview = memref.subview %alloc [%l] [...] [1]
 /// memref.tensor_store %t, %subview
 /// %0 = bufferization.to_tensor %alloc restrict writable
 ///
 /// In addition to rewriting the IR as shown above, this function returns the
-/// newly allocated buffer.
+/// newly allocated buffer. The `insertionPoint` parameter can be used to
+/// specify a custom insertion point for the buffer allocation.
 Value bufferizeToAllocation(RewriterBase &rewriter, tensor::PadOp padOp,
-                            Attribute memorySpace = {});
+                            Attribute memorySpace = {},
+                            Operation *insertionPoint = nullptr);
+
+/// Materialize a buffer allocation for the given vector.mask op and bufferize
+/// the op, including its region. E.g.:
+///
+/// %0 = vector.mask {
+///   vector.transfer_write %v, %t : vector<16xf32>, tensor<?xf32>
+/// } : vector<16xi1> -> tensor<?xf32>
+///
+/// is lowered to:
+///
+/// %alloc = memref.alloc
+/// memref.tensor_store %t, %subview
+/// vector.mask {
+///   vector.transfer_write %arg0, %alloc : vector<16xf32>, memref<?xf32>
+/// } : vector<16xi1>
+/// %0 = bufferization.to_tensor %alloc restrict writable
+///
+/// In addition to rewriting the IR as shown above, this function returns the
+/// newly allocated buffer. The `insertionPoint` parameter can be used to
+/// specify a custom insertion point for the buffer allocation.
+Value bufferizeToAllocation(RewriterBase &rewriter, vector::MaskOp maskOp,
+                            Attribute memorySpace = {},
+                            Operation *insertionPoint = nullptr);
 
 /// Bufferize the given op with tensor semantics and materialize the result in
 /// a newly allocated buffer.
 ///
 /// Only bufferizable ops that bufferize to a memory write or have an
 /// aliasing OpOperand (and do not themselves bufferize to an allocation) are
 /// supported. They are bufferized using their BufferizableOpInterface
 /// implementation.
 ///
-/// Selected ops that bufferize to an allocation are also supported:
+/// Selected ops that bufferize to an allocation (or need special handling) are
+/// also supported:
 /// - tensor.pad
+/// - vector.mask
+///
+/// This function returns the newly allocated buffer. The `insertionPoint`
+/// parameter can be used to specify a custom insertion point for the buffer
+/// allocation.
 Value bufferizeToAllocation(RewriterBase &rewriter, Operation *op,
-                            Attribute memorySpace = {});
+                            Attribute memorySpace = {},
+                            Operation *insertionPoint = nullptr);
 
 /// Try to eliminate tensor::EmptyOps inside `op` that are anchored on a
 /// LinalgOp. This transforms looks for LinalgOps that have an unused output
 /// operand and an input operand that is rooted in a tensor::EmptyOp. The
 /// tensor::EmptyOp uses are replaced with the output operand and the two
 /// operands of the LinalgOp are swapped.
 ///
 /// Example:

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

static Value createAllocationForTensor(RewriterBase &rewriter, Location loc,
   // Place deallocation at the end of the block.
   rewriter.setInsertionPoint(rewriter.getInsertionBlock()->getTerminator());
   rewriter.create<memref::DeallocOp>(loc, alloc);
 
   return alloc;
 }
 
 Value linalg::bufferizeToAllocation(RewriterBase &rewriter, PadOp padOp,
-                                    Attribute memorySpace) {
+                                    Attribute memorySpace,
+                                    Operation *insertionPoint) {
   OpBuilder::InsertionGuard g(rewriter);
-  rewriter.setInsertionPoint(padOp);
+  rewriter.setInsertionPoint(insertionPoint ? insertionPoint : padOp);
   Location loc = padOp.getLoc();
 
   // Create buffer allocation.
   Value alloc =
       createAllocationForTensor(rewriter, loc, padOp.getResult(), memorySpace);
-  rewriter.setInsertionPointAfter(alloc.getDefiningOp());
+  rewriter.setInsertionPoint(padOp);
 
   // Create linalg.fill or linalg.generic.
   Operation *fillOp = movePaddingToFillOrGenericOp(rewriter, loc, padOp, alloc);
   rewriter.setInsertionPointAfter(fillOp);
 
   // Create memref.tensor_store.
   SmallVector<OpFoldResult> sizes =
       getMixedSizes(rewriter, loc, padOp.getSource());
   SmallVector<OpFoldResult> strides(padOp.getResultType().getRank(),
                                     rewriter.getIndexAttr(1));
   Value subview = rewriter.create<memref::SubViewOp>(
       loc, alloc, /*offsets=*/padOp.getMixedLowPad(), sizes, strides);
   rewriter.create<memref::TensorStoreOp>(loc, padOp.getSource(), subview);
 
   // Create bufferization.to_tensor with "restrict" and "writable". The returned
   // tensor is a new buffer allocation, so it does not alias with any buffer.
   Value toTensorOp = rewriter.create<bufferization::ToTensorOp>(
       loc, alloc, /*restrict=*/true, /*writable=*/true);
   rewriter.replaceOp(padOp, toTensorOp);
   return alloc;
 }
 
+Value linalg::bufferizeToAllocation(RewriterBase &rewriter,
+                                    vector::MaskOp maskOp,
+                                    Attribute memorySpace,
+                                    Operation *insertionPoint) {
+  assert(llvm::range_size(maskOp.getMaskBlock()->without_terminator()) == 1 &&
+         "expected single masked op");
+  OpBuilder::InsertionGuard g(rewriter);
+  bufferization::BufferizationOptions options;
+  Operation *yieldOp = maskOp.getMaskRegion().front().getTerminator();
+  assert(isa<vector::YieldOp>(yieldOp) && "expected yield op terminator");
+
+  // Bufferize maskable op. By default, place the buffer allocation right before
+  // the mask op.
+  Value alloc = bufferizeToAllocation(
+      rewriter, maskOp.getMaskableOp(), memorySpace,
+      /*insertionPoint=*/insertionPoint ? insertionPoint : maskOp);
+
+  // Bufferize terminator.
+  rewriter.setInsertionPoint(yieldOp);
+  if (failed(cast<bufferization::BufferizableOpInterface>(yieldOp).bufferize(
+          rewriter, options)))
+    return nullptr;
+
+  // Erase dead to_tensor ops inside of the mask op. This is necessary because
+  // there only be one op (apart from the terminator) inside the mask op.
+  // TODO: Remove dead to_tensor ops more aggressively during bufferization.
+  SmallVector<Operation *> toTensorOps;
+  maskOp.walk([&](bufferization::ToTensorOp toTensorOp) {
+    if (toTensorOp->getUses().empty())
+      toTensorOps.push_back(toTensorOp.getOperation());
+  });
+  for (Operation *op : toTensorOps)
+    rewriter.eraseOp(op);
+
+  // Bufferize mask op.
+  SmallVector<OpOperand *> resultUses;
+  for (Value result : maskOp.getResults())
+    if (isa<TensorType>(result.getType()))
+      for (OpOperand &use : result.getUses())
+        resultUses.push_back(&use);
+  rewriter.setInsertionPoint(maskOp);
+  if (failed(cast<bufferization::BufferizableOpInterface>(maskOp.getOperation())
+                 .bufferize(rewriter, options)))
+    return nullptr;
+
+  // Set "restrict" attribute, indicating that no other tensor aliases with
+  // this tensor. That is because we just allocated a new buffer for the tensor.
+  for (OpOperand *resultUse : resultUses) {
+    auto toTensorOp =
+        resultUse->get().getDefiningOp<bufferization::ToTensorOp>();
+    assert(toTensorOp && "expected to_tensor op");
+    rewriter.updateRootInPlace(toTensorOp, [&]() {
+      toTensorOp.setRestrict(true);
+      toTensorOp.setWritable(true);
+    });
+  }
+
+  return alloc;
+}
+
 /// Lower tensor.from_elements to a sequence of chained tensor.insert.
 FailureOr<Operation *> mlir::linalg::rewriteInDestinationPassingStyle(
     RewriterBase &rewriter, tensor::FromElementsOp fromElementsOp) {
   Location loc = fromElementsOp.getLoc();
   RankedTensorType tensorType =
       cast<RankedTensorType>(fromElementsOp.getType());
   auto shape = tensorType.getShape();
 
SmallVector<OpFoldResult> sliceStrides(resultType.getRank(),
                                          rewriter.getIndexAttr(1));
   auto insertSliceOp = rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
       padOp, padOp.getSource(), fillOp->getResult(0),
       /*offsets=*/padOp.getMixedLowPad(), sliceSizes, sliceStrides);
   return insertSliceOp.getOperation();
 }
 
 Value linalg::bufferizeToAllocation(RewriterBase &rewriter, Operation *op,
-                                    Attribute memorySpace) {
+                                    Attribute memorySpace,
+                                    Operation *insertionPoint) {
   using namespace bufferization;
 
   // Call specialized overload for certain ops.
   if (auto padOp = dyn_cast<tensor::PadOp>(op))
     return bufferizeToAllocation(rewriter, padOp, memorySpace);
+  if (auto maskOp = dyn_cast<vector::MaskOp>(op))
+    return bufferizeToAllocation(rewriter, maskOp, memorySpace);
 
   // Only bufferizable ops are supported.
   auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op);
   if (!bufferizableOp)
     return nullptr;
   BufferizationOptions options;
   AnalysisState state(options);
 
for (OpOperand &operand : op->getOpOperands()) {
     addOutOfPlaceOperand(&operand);
   }
   // TODO: Support multiple buffers.
   if (outOfPlaceOperands.size() != 1)
     return nullptr;
 
   // Allocate buffers.
   OpBuilder::InsertionGuard g(rewriter);
-  rewriter.setInsertionPoint(op);
+  rewriter.setInsertionPoint(insertionPoint ? insertionPoint : op);
   SmallVector<Value> allocs;
   for (OpOperand *operand : outOfPlaceOperands) {
     Value alloc = createAllocationForTensor(rewriter, op->getLoc(),
                                             operand->get(), memorySpace);
     allocs.push_back(alloc);
     // Initialize buffer with a copy of the operand data.
     // TODO: Do not copy uninitialized tensors such as tensor.empty.
     rewriter.create<memref::TensorStoreOp>(op->getLoc(), operand->get(), alloc);
     rewriter.updateRootInPlace(op, [&]() {
       operand->set(rewriter.create<ToTensorOp>(op->getLoc(), alloc));
     });
   }
 
   // Bufferize the op.
+  rewriter.setInsertionPoint(op);
   if (failed(bufferizableOp.bufferize(rewriter, options)))
     return nullptr;
 
   // Set "restrict" attribute, indicating that no other tensor aliases with
   // this tensor. That is because we just allocated a new buffer for the tensor.
   for (OpOperand *resultUse : resultUses) {
     auto toTensorOp = resultUse->get().getDefiningOp<ToTensorOp>();
     assert(toTensorOp && "expected to_tensor op");

mlir/test/Dialect/Linalg/transform-op-bufferize-to-allocation.mlir

 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["tensor.extract"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   // expected-error @below{{failed to bufferize operation}}
   %2 = transform.structured.bufferize_to_allocation %0 {memory_space = 4} : !transform.any_op
 }
+
+// -----
+
+// CHECK-LABEL: func @vector_mask(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?xf32>,
+//       CHECK:   %[[alloc:.*]] = memref.alloc(%{{.*}}) : memref<?xf32, 4>
+//       CHECK:   memref.tensor_store %[[t]], %[[alloc]]
+//       CHECK:   vector.mask %{{.*}} { vector.transfer_write %{{.*}}, %[[alloc]]
+//       CHECK:   %[[r:.*]] = bufferization.to_tensor %[[alloc]] restrict writable
+//       CHECK:   memref.dealloc %[[alloc]]
+//       CHECK:   return %[[r]]
+func.func @vector_mask(%t: tensor<?xf32>, %val: vector<16xf32>, %idx: index, %m0: vector<16xi1>) -> tensor<?xf32> {
+  %r = vector.mask %m0 { vector.transfer_write %val, %t[%idx] : vector<16xf32>, tensor<?xf32> } : vector<16xi1> -> tensor<?xf32>
+  return %r : tensor<?xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["vector.mask"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %2 = transform.structured.bufferize_to_allocation %0 {memory_space = 4} : !transform.any_op
+}