[mlir][vector][bufferize] Bufferize vector.mask and vector.yield

[springerm]

Summary

The masked op can currently not bufferize out-of-place. Such IR would be rejected by the One-Shot Bufferize because it would mean that a new buffer allocation is yielded from a block. Furthermore, only one operation is currently allowed inside vector.mask.

Depends On: D141685

[springerm]

Extract areAllocationsAllowedInRegion into separate revision.

[dcaballe]

mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp

148–150

You should be able to cast here as this terminator is enforced by the op verifier.

159

-> mask.getMaskableOp()

163–188

I'm wondering if we could simplify all this logic just by masking the new buferized operation with a new mask. I'm introducing a maskOperation utility that takes care of creating the new yield with the new return values, etc.

[springerm]

address comments and better error message

[springerm]

mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp

163–188

Yes that would make the code a bit shorter here. Overall, I'm not too worried about it because this is common pattern that we also have with other ops (during bufferization some results change, others do not). Let me know if you add a new helper function, then I will use it here.

[dcaballe]

Thanks!

mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp

LogicalResult bufferization::bufferizeOp(Operation *op,
 
   // Fold all to_memref(to_tensor(x)) pairs.
   for (Operation *op : toMemrefOps) {
     rewriter.setInsertionPoint(op);
     (void)bufferization::foldToMemrefToTensorPair(rewriter,
                                                   cast<ToMemrefOp>(op));
   }
 
+  // Remove all dead to_tensor ops.
+  op->walk<WalkOrder::PostOrder>([&](ToTensorOp toTensorOp) {
+    if (toTensorOp->getUses().empty()) {
+      rewriter.eraseOp(toTensorOp);
+      return WalkResult::skip();
+    }
+    return WalkResult::advance();
+  });
+
   /// Check the result of bufferization. Return an error if an op was not
   /// bufferized, unless partial bufferization is allowed.
   if (options.allowUnknownOps)
     return success();
 
   for (Operation *op : worklist) {
     // Skip ops that are entirely gone.
     if (erasedOps.contains(op))

mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp

 //===- BufferizableOpInterfaceImpl.cpp - Impl. of BufferizableOpInterface -===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.h"
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/IR/DstBufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Operation.h"
 
 using namespace mlir;
 using namespace mlir::bufferization;
 using namespace mlir::vector;
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     replaceOpWithNewBufferizedOp<vector::GatherOp>(
         rewriter, gatherOp, gatherOp.getVectorType(), *buffer,
         gatherOp.getIndices(), gatherOp.getIndexVec(), gatherOp.getMask(),
         gatherOp.getPassThru());
     return success();
   }
 };
 
+/// Bufferization of vector.mask. Replaced with a new vector.mask that
+/// operates on a memref.
+struct MaskOpInterface
+    : public BufferizableOpInterface::ExternalModel<MaskOpInterface,
+                                                    vector::MaskOp> {
+  SmallVector<OpOperand *>
+  getAliasingOpOperand(Operation *op, OpResult opResult,
+                       const AnalysisState &state) const {
+    // MaskOps do not have tensor OpOperands. The yielded values are the result
+    // of the wrapped op.
+    auto maskOp = cast<vector::MaskOp>(op);
+    size_t resultNum = std::distance(op->getOpResults().begin(),
+                                     llvm::find(op->getOpResults(), opResult));
+    auto yieldOp =
+        cast<vector::YieldOp>(maskOp.getMaskRegion().front().getTerminator());
+    return {&yieldOp->getOpOperand(resultNum)};
+  }
+
+  LogicalResult resolveConflicts(Operation *op, RewriterBase &rewriter,
+                                 const AnalysisState &state) const {
+    auto bufferizableOp = cast<BufferizableOpInterface>(op);
+    if (failed(bufferizableOp.resolveTensorOpOperandConflicts(rewriter, state)))
+      return failure();
+
+    // TODO: Remove this function when vector.mask bodies can bufferize
+    // out-of-place. This is currently not supported because yielding allocs
+    // from a block leads to a memory leak and because vector.mask supports only
+    // a single op in its body.
+    auto maskOp = cast<vector::MaskOp>(op);
+    if (!maskOp.getMaskRegion()
+             .front()
+             .getOps<bufferization::AllocTensorOp>()
+             .empty())
+      return op->emitOpError("body must bufferize in-place");
+
+    return success();
+  }
+
+  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+                          const BufferizationOptions &options) const {
+    auto maskOp = cast<vector::MaskOp>(op);
+
+    // Do not bufferize if the masked op is not bufferizable.
+    Operation *maskedOp = maskOp.getMaskableOp();
+    if (!options.dynCastBufferizableOp(maskedOp))
+      return success();
+
+    // Update the terminator: Drop all operands that are not results of the
+    // masked op.
+    auto yieldOp =
+        cast<vector::YieldOp>(maskOp.getMaskRegion().front().getTerminator());
+    SmallVector<Value> newReturnValues(maskOp->getNumResults(), Value());
+    SmallVector<Value> newYieldedValues;
+    for (const auto &it : llvm::enumerate(yieldOp.getOperands())) {
+      if (llvm::find(maskedOp->getOpResults(), it.value()) !=
+          maskedOp->getOpResults().end()) {
+        newYieldedValues.push_back(it.value());
+      } else {
+        // This used to be a tensor result of the masked op, but is now a memref
+        // that is defined outside of the vector.mask op.
+        newReturnValues[it.index()] = it.value();
+      }
+    }
+    rewriter.updateRootInPlace(yieldOp, [&]() {
+      yieldOp.getOperandsMutable().assign(newYieldedValues);
+    });
+
+    // Create a new vector.mask op.
+    TypeRange newResultTypes(newYieldedValues);
+    auto newOp = rewriter.create<vector::MaskOp>(
+        op->getLoc(), newResultTypes, maskOp.getMask(), maskOp.getPassthru(),
+        /*maskableOp=*/nullptr,
+        /*maskRegionBuilder=*/[](OpBuilder &b, Operation *) {});
+    newOp.getRegion().takeBody(maskOp.getMaskRegion());
+
+    // Replace all uses of the old vector.mask op.
+    int idx = 0;
+    for (int i = 0; i < static_cast<int>(maskOp->getNumResults()); ++i) {
+      if (!newReturnValues[i])
+        newReturnValues[i] = newOp->getResult(idx++);
+    }
+    replaceOpWithBufferizedValues(rewriter, maskOp, newReturnValues);
+    return success();
+  }
+
+  BufferRelation bufferRelation(Operation *op, OpResult opResult,
+                                const AnalysisState &state) const {
+    return BufferRelation::Equivalent;
+  }
+};
+
+/// Bufferization of vector.yield. Replaced with a new vector.yield that
+/// operates on a memref.
+struct YieldOpInterface
+    : public BufferizableOpInterface::ExternalModel<YieldOpInterface,
+                                                    vector::YieldOp> {
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const AnalysisState &state) const {
+    return true;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const AnalysisState &state) const {
+    return false;
+  }
+
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
+    return {op->getParentOp()->getResult(opOperand.getOperandNumber())};
+  }
+
+  bool mustBufferizeInPlace(Operation *op, OpOperand &opOperand,
+                            const AnalysisState &state) const {
+    // Yield operands always bufferize inplace. Otherwise, an alloc + copy
+    // may be generated inside the block. We should not return/yield allocations
+    // when possible.
+    return true;
+  }
+
+  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+                          const BufferizationOptions &options) const {
+    auto yieldOp = cast<vector::YieldOp>(op);
+
+    // Only supported as a vector.mask terminator.
+    auto maskOp = dyn_cast<vector::MaskOp>(yieldOp->getParentOp());
+    if (!maskOp)
+      return yieldOp->emitError("unsupported vector::YieldOp parent");
+
+    // Do not bufferize if the masked op is not bufferizable.
+    Operation *maskedOp = &maskOp.getMaskRegion().front().front();
+    if (!options.dynCastBufferizableOp(maskedOp))
+      return success();
+
+    // Create a new terminator with the same number of operands. Some of these
+    // may get dropped during the bufferization of vector.mask.
+    SmallVector<Value> newResults;
+    for (Value value : yieldOp.getOperands()) {
+      if (value.getType().isa<TensorType>()) {
+        FailureOr<Value> maybeBuffer = getBuffer(rewriter, value, options);
+        if (failed(maybeBuffer))
+          return failure();
+        newResults.push_back(*maybeBuffer);
+      } else {
+        newResults.push_back(value);
+      }
+    }
+
+    replaceOpWithNewBufferizedOp<vector::YieldOp>(rewriter, op, newResults);
+    return success();
+  }
+};
+
 } // namespace
 } // namespace vector
 } // namespace mlir
 
 void mlir::vector::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, vector::VectorDialect *dialect) {
     TransferReadOp::attachInterface<TransferReadOpInterface>(*ctx);
     TransferWriteOp::attachInterface<TransferWriteOpInterface>(*ctx);
     GatherOp::attachInterface<GatherOpInterface>(*ctx);
+    MaskOp::attachInterface<MaskOpInterface>(*ctx);
+    YieldOp::attachInterface<YieldOpInterface>(*ctx);
   });
 }

mlir/test/Dialect/Vector/bufferize-invalid.mlir

+// RUN: mlir-opt %s -vector-bufferize -split-input-file -verify-diagnostics
+// | FileCheck %s
+
+// CHECK-LABEL: func @mask(
+func.func @mask(%t0: tensor<?xf32>, %val: vector<16xf32>, %idx: index, %m0: vector<16xi1>) -> tensor<?xf32> {
+  // expected-error @+1 {{'vector.mask' op body must bufferize in-place}}
+  %0 = vector.mask %m0 { vector.transfer_write %val, %t0[%idx] : vector<16xf32>, tensor<?xf32> } : vector<16xi1> -> tensor<?xf32>
+  return %0 : tensor<?xf32>
+}

mlir/test/Dialect/Vector/bufferize.mlir

 //       CHECK:   %[[m:.*]] = bufferization.to_memref %[[base]] : memref<?x?xf32>
 //       CHECK:   %[[c0:.*]] = arith.constant 0 : index
 //       CHECK:   %[[out:.*]] = vector.gather %[[m]][%[[c0]], %[[c0]]] [%[[v]]], %[[mask]], %[[pass_thru]] : memref<?x?xf32>, vector<16xi32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
 func.func @gather(%base: tensor<?x?xf32>, %v: vector<16xi32>, %mask: vector<16xi1>, %pass_thru: vector<16xf32>) -> vector<16xf32> {
   %c0 = arith.constant 0 : index
   %0 = vector.gather %base[%c0, %c0][%v], %mask, %pass_thru : tensor<?x?xf32>, vector<16xi32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
   return %0 : vector<16xf32>
 }
+
+// TODO: Add test case for vector.mask. The masked op can currently not
+// bufferize out-of-place, so the only test case is in one-shot-bufferize.mlir.

mlir/test/Dialect/Vector/one-shot-bufferize.mlir

+// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries" -split-input-file | FileCheck %s
+
+// CHECK-LABEL: func @mask(
+//  CHECK-SAME:     %[[t0:.*]]: memref<?xf32, strided<[?], offset: ?>>
+func.func @mask(%t0: tensor<?xf32>, %val: vector<16xf32>, %idx: index, %m0: vector<16xi1>) -> tensor<?xf32> {
+  // CHECK-NOT: alloc
+  // CHECK-NOT: copy
+  //     CHECK: vector.mask %{{.*}} { vector.transfer_write %{{.*}}, %[[t0]][%{{.*}}] : vector<16xf32>, memref<?xf32, strided<[?], offset: ?>> } : vector<16xi1>
+  %0 = vector.mask %m0 { vector.transfer_write %val, %t0[%idx] : vector<16xf32>, tensor<?xf32> } : vector<16xi1> -> tensor<?xf32>
+  //     CHECK: return %[[t0]]
+  return %0 : tensor<?xf32>
+}