[mlir] VectorToSCF target rank is a pass option

[springerm]

Summary

Make "target rank" a pass option of VectorToSCF.

Depends On D102101

[springerm]

no change

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

mlir/include/mlir/Conversion/Passes.td

def ConvertVectorToSCF : FunctionPass<"convert-vector-to-scf"> {
   let dependentDialects = [
     "AffineDialect",
     "memref::MemRefDialect",
     "scf::SCFDialect"
   ];
   let options = [
     Option<"fullUnroll", "full-unroll", "bool", /*default=*/"false",
            "Perform full unrolling when converting vector transfers to SCF">,
+    Option<"targetRank", "target-rank", "unsigned", /*default=*/"1",
+           "Target vector rank to which transfer ops should be lowered">,
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // VectorToLLVM
 //===----------------------------------------------------------------------===//
 
 def ConvertVectorToLLVM : Pass<"convert-vector-to-llvm", "ModuleOp"> {

mlir/include/mlir/Conversion/VectorToSCF/VectorToSCF.h

 ///          vector<17x15xf32>, memref<?x?x?xf32>
 /// ```
 ///
 /// When applying the pattern a second time, the existing alloca() operation
 /// is reused and only a second vector.type_cast is added.
 
 struct VectorTransferToSCFOptions {
   bool unroll = false;
+  unsigned targetRank = 1;
+
   VectorTransferToSCFOptions &setUnroll(bool u) {
     unroll = u;
     return *this;
   }
+
+  VectorTransferToSCFOptions &setTargetRank(unsigned r) {
+    targetRank = r;
+    return *this;
+  }
 };
 
 /// Collect a set of patterns to convert from the Vector dialect to SCF + std.
 void populateVectorToSCFConversionPatterns(
     RewritePatternSet &patterns,
     const VectorTransferToSCFOptions &options = VectorTransferToSCFOptions());
 
 /// Create a pass to convert a subset of vector ops to SCF.
 std::unique_ptr<Pass> createConvertVectorToSCFPass(
     const VectorTransferToSCFOptions &options = VectorTransferToSCFOptions());
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_VECTORTOSCF_VECTORTOSCF_H_

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

 using vector::TransferReadOp;
 using vector::TransferWriteOp;
 
 namespace {
 
 /// Attribute name used for labeling transfer ops during progressive lowering.
 static const char kPassLabel[] = "__vector_to_scf_lowering__";
 
-/// Lower to 1D transfer ops. Target-specific lowering will lower those.
-static const int64_t kTargetRank = 1;
+/// Patterns that inherit from this struct have access to
+/// VectorTransferToSCFOptions.
+template <typename OpTy>
+struct VectorToSCFPattern : public OpRewritePattern<OpTy> {
+  explicit VectorToSCFPattern(MLIRContext *context,
+                              VectorTransferToSCFOptions opt)
+      : OpRewritePattern<OpTy>(context), options(opt) {}
+
+  VectorTransferToSCFOptions options;
+};
 
 /// Given a MemRefType with VectorType element type, unpack one dimension from
 /// the VectorType into the MemRefType.
 ///
 /// E.g.: memref<9xvector<5x6xf32>> --> memref<9x5xvector<6xf32>>
 static MemRefType unpackOneDim(MemRefType type) {
   auto vectorType = type.getElementType().dyn_cast<VectorType>();
   auto memrefShape = type.getShape();
static ArrayAttr dropFirstElem(OpBuilder &builder, ArrayAttr attr) {
   if (!attr)
     return attr;
   return ArrayAttr::get(builder.getContext(), attr.getValue().drop_front());
 }
 
 /// Add the pass label to a vector transfer op if its rank is not the target
 /// rank.
 template <typename OpTy>
-static void maybeApplyPassLabel(OpBuilder &builder, OpTy newXferOp) {
-  if (newXferOp.getVectorType().getRank() > kTargetRank)
+static void maybeApplyPassLabel(OpBuilder &builder, OpTy newXferOp,
+                                unsigned targetRank) {
+  if (newXferOp.getVectorType().getRank() > targetRank)
     newXferOp->setAttr(kPassLabel, builder.getUnitAttr());
 }
 
 /// Given a transfer op, find the memref from which the mask is loaded. This
 /// is similar to Strategy<TransferWriteOp>::getBuffer.
 template <typename OpTy>
 static Value getMaskBuffer(OpTy xferOp) {
   assert(xferOp.mask() && "Expected that transfer op has mask");
struct Strategy<TransferReadOp> {
   ///       : memref<?x?x?xf32>, vector<3xf32>
   ///   memref.store %vec, %casted[%i, %j] : memref<5x4xvector<3xf32>>
   /// }
   /// ```
   ///
   /// Note: The loop and type cast are generated in TransferOpConversion.
   ///       The original TransferReadOp and store op are deleted in `cleanup`.
   /// Note: The `mask` operand is set in TransferOpConversion.
-  static TransferReadOp rewriteOp(OpBuilder &builder, TransferReadOp xferOp,
-                                  Value buffer, Value iv) {
+  static TransferReadOp rewriteOp(OpBuilder &builder,
+                                  VectorTransferToSCFOptions options,
+                                  TransferReadOp xferOp, Value buffer,
+                                  Value iv) {
     SmallVector<Value, 8> storeIndices;
     getBufferIndices(xferOp, storeIndices);
     storeIndices.push_back(iv);
 
     SmallVector<Value, 8> xferIndices;
     getXferIndices(xferOp, iv, xferIndices);
 
     auto bufferType = buffer.getType().dyn_cast<ShapedType>();
     auto vecType = bufferType.getElementType().dyn_cast<VectorType>();
     auto inBoundsAttr = dropFirstElem(builder, xferOp.in_boundsAttr());
     auto newXfer =
         vector_transfer_read(
             vecType, xferOp.source(), xferIndices,
             AffineMapAttr::get(unpackedPermutationMap(xferOp, builder)),
             xferOp.padding(), Value(), inBoundsAttr)
             .value;
 
     maybeApplyPassLabel(builder,
-                        dyn_cast<TransferReadOp>(newXfer.getDefiningOp()));
+                        dyn_cast<TransferReadOp>(newXfer.getDefiningOp()),
+                        options.targetRank);
 
     memref_store(newXfer, buffer, storeIndices);
     return newXfer.getDefiningOp<TransferReadOp>();
   }
 
   /// Handle out-of-bounds accesses on the to-be-unpacked dimension: Write
   /// padding value to the temporary buffer.
   static void handleOutOfBoundsDim(OpBuilder & /*builder*/,
struct Strategy<TransferWriteOp> {
   /// accesses on the to-be-unpacked dimension.
   ///
   /// 1. Load an (N-1)-d vector from the (already `vector.type_cast`ed) buffer,
   ///    using the loop iteration variable `iv`.
   /// 2. Generate a new (N-1)-d TransferWriteOp, writing the loaded vector back
   ///    to memory.
   ///
   /// Note: For more details, see comments on Strategy<TransferReadOp>.
-  static TransferWriteOp rewriteOp(OpBuilder &builder, TransferWriteOp xferOp,
-                                   Value buffer, Value iv) {
+  static TransferWriteOp rewriteOp(OpBuilder &builder,
+                                   VectorTransferToSCFOptions options,
+                                   TransferWriteOp xferOp, Value buffer,
+                                   Value iv) {
     SmallVector<Value, 8> loadIndices;
     getBufferIndices(xferOp, loadIndices);
     loadIndices.push_back(iv);
 
     SmallVector<Value, 8> xferIndices;
     getXferIndices(xferOp, iv, xferIndices);
 
     auto vec = memref_load(buffer, loadIndices);
     auto inBoundsAttr = dropFirstElem(builder, xferOp.in_boundsAttr());
     auto newXfer = vector_transfer_write(
         Type(), vec, xferOp.source(), xferIndices,
         AffineMapAttr::get(unpackedPermutationMap(xferOp, builder)), Value(),
         inBoundsAttr);
 
-    maybeApplyPassLabel(builder, newXfer.op);
+    maybeApplyPassLabel(builder, newXfer.op, options.targetRank);
 
     return newXfer;
   }
 
   /// Handle out-of-bounds accesses on the to-be-unpacked dimension.
   static void handleOutOfBoundsDim(OpBuilder &builder, TransferWriteOp xferOp,
                                    Value buffer, Value iv) {}
 
   /// Cleanup after rewriting the op.
   static void cleanup(PatternRewriter &rewriter, TransferWriteOp xferOp) {
     rewriter.eraseOp(xferOp);
   }
 };
 
 template <typename OpTy>
-LogicalResult checkPrepareXferOp(OpTy xferOp) {
+LogicalResult checkPrepareXferOp(OpTy xferOp, unsigned targetRank) {
   if (xferOp->hasAttr(kPassLabel))
     return failure();
-  if (xferOp.getVectorType().getRank() <= kTargetRank)
+  if (xferOp.getVectorType().getRank() <= targetRank)
     return failure();
   return success();
 }
 
 /// Prepare a TransferReadOp for progressive lowering.
 ///
 /// 1. Allocate a temporary buffer.
 /// 2. Label the TransferReadOp, marking it eligible for progressive lowering.

 /// %0 = memref.alloca() : memref<vector<5x4xf32>>
 /// %1 = vector.transfer_read %A[%a, %b, %c], %cst
 ///     { __vector_to_scf_lowering__ } : vector<5x4xf32>, memref<?x?x?xf32>
 /// memref.store %1, %0[] : memref<vector<5x4xf32>>
 /// %vec = memref.load %0[] : memref<vector<5x4xf32>>
 /// ```
 ///
 /// Note: A second temporary buffer may be allocated for the `mask` operand.
-struct PrepareTransferReadConversion : public OpRewritePattern<TransferReadOp> {
-  using OpRewritePattern<TransferReadOp>::OpRewritePattern;
+struct PrepareTransferReadConversion
+    : public VectorToSCFPattern<TransferReadOp> {
+  using VectorToSCFPattern<TransferReadOp>::VectorToSCFPattern;
 
   LogicalResult matchAndRewrite(TransferReadOp xferOp,
                                 PatternRewriter &rewriter) const override {
-    if (checkPrepareXferOp(xferOp).failed())
+    if (checkPrepareXferOp(xferOp, options.targetRank).failed())
       return failure();
 
     ScopedContext scope(rewriter, xferOp.getLoc());
     auto buffers = allocBuffers(xferOp);
     auto *newXfer = rewriter.clone(*xferOp.getOperation());
     newXfer->setAttr(kPassLabel, rewriter.getUnitAttr());
     if (xferOp.mask()) {
       dyn_cast<TransferReadOp>(newXfer).maskMutable().assign(

 /// memref.store %vec, %0[] : memref<vector<5x4xf32>>
 /// %1 = memref.load %0[] : memref<vector<5x4xf32>>
 /// vector.transfer_write %1, %A[%a, %b, %c] { __vector_to_scf_lowering__ }
 ///     : vector<5x4xf32>, memref<?x?x?xf32>
 /// ```
 ///
 /// Note: A second temporary buffer may be allocated for the `mask` operand.
 struct PrepareTransferWriteConversion
-    : public OpRewritePattern<TransferWriteOp> {
-  using OpRewritePattern<TransferWriteOp>::OpRewritePattern;
+    : public VectorToSCFPattern<TransferWriteOp> {
+  using VectorToSCFPattern<TransferWriteOp>::VectorToSCFPattern;
 
   LogicalResult matchAndRewrite(TransferWriteOp xferOp,
                                 PatternRewriter &rewriter) const override {
-    if (checkPrepareXferOp(xferOp).failed())
+    if (checkPrepareXferOp(xferOp, options.targetRank).failed())
       return failure();
 
     ScopedContext scope(rewriter, xferOp.getLoc());
     auto buffers = allocBuffers(xferOp);
     memref_store(xferOp.vector(), buffers.dataBuffer);
     auto loadedVec = memref_load(buffers.dataBuffer);
     rewriter.updateRootInPlace(xferOp, [&]() {
       xferOp.vectorMutable().assign(loadedVec);

 ///    %casted = vector.type_cast %0
 ///        : memref<5xvector<4x3xf32>> to memref<5x4xvector<3xf32>>
 ///    ```
 /// 2. Generate a for loop and rewrite the transfer op according to the
 ///    corresponding Strategy<OpTy>. If the to-be-unpacked dimension can be
 ///    out-of-bounds, generate an if-check and handle both cases separately.
 /// 3. Clean up according to the corresponding Strategy<OpTy>.
 template <typename OpTy>
-struct TransferOpConversion : public OpRewritePattern<OpTy> {
-  using OpRewritePattern<OpTy>::OpRewritePattern;
+struct TransferOpConversion : public VectorToSCFPattern<OpTy> {
+  using VectorToSCFPattern<OpTy>::VectorToSCFPattern;
 
   LogicalResult matchAndRewrite(OpTy xferOp,
                                 PatternRewriter &rewriter) const override {
     if (!xferOp->hasAttr(kPassLabel))
       return failure();
 
     ScopedContext scope(rewriter, xferOp.getLoc());
 
rewriter.create<scf::ForOp>(
         xferOp.getLoc(), lb, ub, step, ValueRange(),
         [&](OpBuilder &b, Location loc, Value iv, ValueRange /*loopState*/) {
           ScopedContext scope(b, loc);
           generateInBoundsCheck(
               xferOp, iv, b, unpackedDim(xferOp),
               /*inBoundsCase=*/
               [&](OpBuilder &b, Location /*loc*/) {
                 // Create new transfer op.
-                OpTy newXfer =
-                    Strategy<OpTy>::rewriteOp(b, xferOp, castedDataBuffer, iv);
+                OpTy newXfer = Strategy<OpTy>::rewriteOp(
+                    b, this->options, xferOp, castedDataBuffer, iv);
 
                 // If old transfer op has a mask: Set mask on new transfer op.
                 // Special case: If the mask of the old transfer op is 1D and
                 // the
                 //               unpacked dim is not a broadcast, no mask is
                 //               needed on the new transfer op.
                 if (xferOp.mask() && (xferOp.isBroadcastDim(0) ||
                                       xferOp.getMaskType().getRank() > 1)) {

 /// %tmp4 = vector.transfer_read %A[%a, %b + 4, %c], %padding
 ///     : memref<?x?x?xf32>, vector<4xf32>
 /// %vec = vector.insert %tmp1, %v3[4] : vector<4xf32> into vector<5x4xf32>
 /// ```
 ///
 /// Note: As an optimization, if the result of the original TransferReadOp
 /// was directly inserted into another vector, no new %v_init vector is created.
 /// Instead, the new TransferReadOp results are inserted into that vector.
-struct UnrollTransferReadConversion : public OpRewritePattern<TransferReadOp> {
-  using OpRewritePattern<TransferReadOp>::OpRewritePattern;
+struct UnrollTransferReadConversion
+    : public VectorToSCFPattern<TransferReadOp> {
+  using VectorToSCFPattern<TransferReadOp>::VectorToSCFPattern;
 
   /// Return the vector into which the newly created TransferReadOp results
   /// are inserted.
   Value getResultVector(TransferReadOp xferOp,
                         PatternRewriter &rewriter) const {
     if (auto insertOp = getInsertOp(xferOp))
       return insertOp.dest();
     return std_splat(xferOp.getVectorType(), xferOp.padding()).value;
if (auto insertOp = getInsertOp(xferOp)) {
       });
     }
   }
 
   /// Rewrite the op: Unpack one dimension. Can handle masks, out-of-bounds
   /// accesses, and broadcasts and transposes in permutation maps.
   LogicalResult matchAndRewrite(TransferReadOp xferOp,
                                 PatternRewriter &rewriter) const override {
-    if (xferOp.getVectorType().getRank() <= kTargetRank)
+    if (xferOp.getVectorType().getRank() <= options.targetRank)
       return failure();
 
     ScopedContext scope(rewriter, xferOp.getLoc());
     auto insertOp = getInsertOp(xferOp);
     auto vec = getResultVector(xferOp, rewriter);
     auto vecType = vec.getType().dyn_cast<VectorType>();
     auto xferVecType = xferOp.getVectorType();
     auto newXferVecType = VectorType::get(xferVecType.getShape().drop_front(),

 /// ```
 ///
 /// Note: As an optimization, if the vector of the original TransferWriteOp
 /// was directly extracted from another vector via an ExtractOp `a`, extract
 /// the vectors for the newly generated TransferWriteOps from `a`'s input. By
 /// doing so, `a` may become dead, and the number of ExtractOps generated during
 /// recursive application of this pattern will be minimal.
 struct UnrollTransferWriteConversion
-    : public OpRewritePattern<TransferWriteOp> {
-  using OpRewritePattern<TransferWriteOp>::OpRewritePattern;
+    : public VectorToSCFPattern<TransferWriteOp> {
+  using VectorToSCFPattern<TransferWriteOp>::VectorToSCFPattern;
 
   /// Return the vector from which newly generated ExtracOps will extract.
   Value getDataVector(TransferWriteOp xferOp) const {
     if (auto extractOp = getExtractOp(xferOp))
       return extractOp.vector();
     return xferOp.vector();
   }
 
if (auto extractOp = getExtractOp(xferOp)) {
       });
     }
   }
 
   /// Rewrite the op: Unpack one dimension. Can handle masks, out-of-bounds
   /// accesses, and broadcasts and transposes in permutation maps.
   LogicalResult matchAndRewrite(TransferWriteOp xferOp,
                                 PatternRewriter &rewriter) const override {
-    if (xferOp.getVectorType().getRank() <= kTargetRank)
+    if (xferOp.getVectorType().getRank() <= options.targetRank)
       return failure();
 
     ScopedContext scope(rewriter, xferOp.getLoc());
     auto vec = getDataVector(xferOp);
     auto xferVecType = xferOp.getVectorType();
     int64_t dimSize = xferVecType.getShape()[0];
 
     // Generate fully unrolled loop of transfer ops.

 /// Is rewritten to approximately the following pseudo-IR:
 /// ```
 /// for i = 0 to 9 {
 ///   %t = vector.extractelement %vec[i] : vector<9xf32>
 ///   memref.store %t, %arg0[%a + i, %b] : memref<?x?xf32>
 /// }
 /// ```
 template <typename OpTy>
-struct TransferOp1dConversion : public OpRewritePattern<OpTy> {
-  using OpRewritePattern<OpTy>::OpRewritePattern;
+struct TransferOp1dConversion : public VectorToSCFPattern<OpTy> {
+  using VectorToSCFPattern<OpTy>::VectorToSCFPattern;
 
   LogicalResult matchAndRewrite(OpTy xferOp,
                                 PatternRewriter &rewriter) const override {
     ScopedContext scope(rewriter, xferOp.getLoc());
     auto map = xferOp.permutation_map();
     auto memRefType = xferOp.getShapedType().template dyn_cast<MemRefType>();
 
     if (!memRefType)

 } // namespace
 
 namespace mlir {
 
 void populateVectorToSCFConversionPatterns(
     RewritePatternSet &patterns, const VectorTransferToSCFOptions &options) {
   if (options.unroll) {
     patterns.add<UnrollTransferReadConversion, UnrollTransferWriteConversion>(
-        patterns.getContext());
+        patterns.getContext(), options);
   } else {
     patterns.add<PrepareTransferReadConversion, PrepareTransferWriteConversion,
                  TransferOpConversion<TransferReadOp>,
-                 TransferOpConversion<TransferWriteOp>>(patterns.getContext());
+                 TransferOpConversion<TransferWriteOp>>(patterns.getContext(),
+                                                        options);
   }
 
-  if (kTargetRank == 1) {
+  if (options.targetRank == 1) {
     patterns.add<TransferOp1dConversion<TransferReadOp>,
-                 TransferOp1dConversion<TransferWriteOp>>(
-        patterns.getContext());
+                 TransferOp1dConversion<TransferWriteOp>>(patterns.getContext(),
+                                                          options);
   }
 }
 
 } // namespace mlir
 
 namespace {
 
 struct ConvertVectorToSCFPass
     : public ConvertVectorToSCFBase<ConvertVectorToSCFPass> {
   ConvertVectorToSCFPass() = default;
   ConvertVectorToSCFPass(const VectorTransferToSCFOptions &options) {
     this->fullUnroll = options.unroll;
+    this->targetRank = options.targetRank;
   }
 
   void runOnFunction() override {
+    VectorTransferToSCFOptions options;
+    options.setUnroll(fullUnroll);
+    options.setTargetRank(targetRank);
+
     RewritePatternSet patterns(getFunction().getContext());
-    populateVectorToSCFConversionPatterns(
-        patterns, VectorTransferToSCFOptions().setUnroll(fullUnroll));
+    populateVectorToSCFConversionPatterns(patterns, options);
     (void)applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
   }
 };
 
 } // namespace
 
 std::unique_ptr<Pass>
 mlir::createConvertVectorToSCFPass(const VectorTransferToSCFOptions &options) {
   return std::make_unique<ConvertVectorToSCFPass>(options);
 }