[mlir][linalg] Add support for dynamic pad op masking vectorization.

[hanchung]

clang-format

[awarzynski]

Thanks for working on this @hanchung ! Just a few minor comments inline.

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

1330

Could you add a test to exercise this?

1473

[nit] I appreciate that this is unrelated to this change, but "Input vector sizes" is a bit misleading IMHO. "Requested vector sizes" or "Target vector sizes" would be more descriptive.

IIUC, inputVectorSizes is an input to vectorize, i.e. the vectorizer. It's just too easy to confuse this with with "input to the PadOp that's being analyzed here". Perhaps that's just me.

1475–1476

This isn't required for this change, is it? I don't mind, but if that's the case then could you add a note in the commit msg?

Also, I'd move this to the very top of this method (it's one of they key assumptions, isn't it?).

1483

[nit] Same note about "input" as above - perhaps "targert"?

[dcaballe]

Awesome! Just a few comments. LGTM overall

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

1330

Adding some comments here would also help understanding.

I understand that this condition is only met for dynamic cases because the padOp semantics prevent the use of any other types that mismatch, right?

1472–1475

I think all this code is shared with vectorize. Is it possible to refactor it to a utility function?

[nicolasvasilache]

Cool, thanks for extending to also support this case.

LGTM once comments are addressed.

[hanchung]

address comments

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

1330

This is because the pad op is not a DPS operation, we don't have a destination tensor to do transfer_write. We should just create the empty op with given low, high and source shape.

1473

I agree with the point, it sounds really reasonable to me. There are many other places using input vector terms in this file, and I'd like to keep the consistency. This should be done in a separate patch, IMO.

1475–1476

I think this is required for all masking vectorization? My understanding is that we need the static value to create the mask. This is the same check in generic op version. I refactored the common logic to a util, please take a look.

[awarzynski]

Thank you!

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

1475–1476

I think this is required for all masking vectorization?

Yes, that was my point :) Your patch is about "dynamic pad op masking vectorization" and this change applies to everything. Like I said, I don't mind :)

[dcaballe]

LGTM, thanks! Please, wait for @awarzynski approval before landing.

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

1363

doc

[dcaballe]

Oh, my review page needed a refresh :)

[hanchung]

add doc

[hanchung]

rebase to main

[hanchung]

fixes for transform ops changes

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

vectorizeAsTensorPadOp(RewriterBase &rewriter, tensor::PadOp padOp,
   Location loc = padOp.getLoc();
   int64_t rank = inputVectorSizes.size();
   auto maskType = VectorType::get(inputVectorSizes, rewriter.getI1Type());
   auto vectorType = VectorType::get(inputVectorSizes, padValue.getType());
 
   // transfer_write_in_bounds(transfer_read_masked(pad_source, pad_value))
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPoint(padOp);
-  auto zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
-  auto emptyOp =
-      rewriter.create<tensor::EmptyOp>(loc, padOp.getResultType(),
-                                       /*dynamicSizes=*/ValueRange{});
+
+  ReifiedRankedShapedTypeDims reifiedReturnShapes;
+  LogicalResult status =
+      cast<ReifyRankedShapedTypeOpInterface>(padOp.getOperation())
+          .reifyResultShapes(rewriter, reifiedReturnShapes);
+  assert(succeeded(status) && "failed to reify result shapes");
+  auto emptyOp = rewriter.create<tensor::EmptyOp>(loc, reifiedReturnShapes[0],
+                                                  padValue.getType());
   SmallVector<OpFoldResult> mixedSourceDims =
       getMixedDimensions(rewriter, loc, padOp.getSource());
   Value mask =
       rewriter.create<vector::CreateMaskOp>(loc, maskType, mixedSourceDims);
+  auto zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
   auto transferReadOp = rewriter.create<vector::TransferReadOp>(
       loc,
       /*vectorType=*/vectorType,
       /*source=*/padOp.getSource(),
       /*indices=*/SmallVector<Value>(rank, zero),
       /*padding=*/padValue,
       /*inBounds=*/SmallVector<bool>(rank, true));
   auto maskedOp = cast<vector::MaskOp>(
       mlir::vector::maskOperation(rewriter, transferReadOp, mask));
-  auto transferWriteOp = rewriter.create<vector::TransferWriteOp>(
+  Operation *write = rewriter.create<vector::TransferWriteOp>(
       loc,
       /*vector=*/maskedOp->getResult(0),
       /*source=*/emptyOp,
       /*indices=*/SmallVector<Value>(rank, zero),
       /*inBounds=*/SmallVector<bool>(rank, true));
-  newResults.push_back(transferWriteOp.getResult());
+  write = mlir::vector::maskOperation(rewriter, write, mask);
+  newResults.push_back(write->getResult(0));
   return success();
 }
 
 // TODO: probably need some extra checks for reduction followed by consumer
 // ops that may not commute (e.g. linear reduction + non-linear instructions).
 static LogicalResult reductionPreconditions(LinalgOp op) {
   if (llvm::none_of(op.getIteratorTypesArray(), isReductionIterator)) {
     LDBG("reduction precondition failed: no reduction iterator\n");
     return failure();
   }
static LogicalResult vectorizeDynamicLinalgOpPrecondition(linalg::LinalgOp op) {
   if (!isa<linalg::GenericOp, linalg::FillOp, linalg::CopyOp,
            linalg::ContractionOpInterface>(op.getOperation()))
     return failure();
 
   LDBG("Dynamically-shaped op meets vectorization pre-conditions\n");
   return success();
 }
 
+/// Returns success if `inputVectorSizes` is a valid masking configuraion for
+/// given `shape`, i.e., it meets:
+///   1. The numbers of elements in both array are equal.
+///   2. `inputVectorSizes` does nos have dynamic dimensions.
+///   3. All the values in `inputVectorSizes` are greater than or equal to
+///      static sizes in `shape`.
 static LogicalResult
-vectorizeLinalgOpPrecondition(LinalgOp linalgOp,
-                              ArrayRef<int64_t> inputVectorSizes,
-                              bool vectorizeNDExtract) {
-  // tensor with dimension of 0 cannot be vectorized.
-  if (llvm::any_of(linalgOp.getStaticShape(),
-                   [](int64_t dim) { return dim == 0; }))
+isValidMaskedInputVector(ArrayRef<int64_t> shape,
+                         ArrayRef<int64_t> inputVectorSizes) {
+  if (inputVectorSizes.size() != shape.size()) {
+    LDBG("Input vector sizes don't match the number of loops");
     return failure();
-  // Check API contract for input vector sizes.
-  if (!inputVectorSizes.empty()) {
-    assert(inputVectorSizes.size() == linalgOp.getNumLoops() &&
-           "Input vector sizes don't match the number of loops");
-    assert(!ShapedType::isDynamicShape(inputVectorSizes) &&
-           "Input vector sizes can't have dynamic dimensions");
-    assert(
-        llvm::all_of(
-            llvm::zip(linalgOp.getStaticLoopRanges(), inputVectorSizes),
+  }
+  if (ShapedType::isDynamicShape(inputVectorSizes)) {
+    LDBG("Input vector sizes can't have dynamic dimensions");
+    return failure();
+  }
+  if (!llvm::all_of(llvm::zip(shape, inputVectorSizes),
                     [](std::tuple<int64_t, int64_t> sizePair) {
                       int64_t staticSize = std::get<0>(sizePair);
                       int64_t inputSize = std::get<1>(sizePair);
                       return ShapedType::isDynamic(staticSize) ||
                              staticSize <= inputSize;
-            }) &&
-        "Input vector sizes must be greater than or equal to iteration space "
+                    })) {
+    LDBG("Input vector sizes must be greater than or equal to iteration space "
          "static sizes");
-  }
+    return failure();
+  }
+  return success();
 
+static LogicalResult
+vectorizeLinalgOpPrecondition(LinalgOp linalgOp,
+                              ArrayRef<int64_t> inputVectorSizes,
+                              bool vectorizeNDExtract) {
+  // tensor with dimension of 0 cannot be vectorized.
+  if (llvm::any_of(linalgOp.getStaticShape(),
+                   [](int64_t dim) { return dim == 0; }))
+    return failure();
+  // Check API contract for input vector sizes.
+  if (!inputVectorSizes.empty() &&
+      failed(isValidMaskedInputVector(linalgOp.getStaticLoopRanges(),
+                                      inputVectorSizes)))
+    return failure();
+
   if (linalgOp.hasDynamicShape() &&
       failed(vectorizeDynamicLinalgOpPrecondition(linalgOp))) {
     LDBG("Dynamically-shaped op failed vectorization pre-conditions\n");
     return failure();
   }
 
   SmallVector<CustomVectorizationPrecondition> customPreconditions;
 
vectorizePadOpPrecondition(tensor::PadOp padOp,
                            ArrayRef<int64_t> inputVectorSizes) {
   auto padValue = padOp.getConstantPaddingValue();
   if (!padValue) {
     LDBG("pad value is not constant: " << padOp << "\n");
     return failure();
   }
 
   ArrayRef<int64_t> resultTensorShape = padOp.getResultType().getShape();
-  if (!(resultTensorShape == inputVectorSizes)) {
-    LDBG("result tensor shape must match input vector sizes: " << padOp
-                                                               << "\n");
+  if (failed(isValidMaskedInputVector(resultTensorShape, inputVectorSizes)))
     return failure();
-  }
 
   if (llvm::any_of(padOp.getLow(), [](Value v) {
         std::optional<int64_t> res = getConstantIntValue(v);
         return !res.has_value() || res.value() != 0;
       })) {
     LDBG("low pad must all be zero: " << padOp << "\n");
     return failure();
   }
 
   return success();
 }
 
 LogicalResult
 mlir::linalg::vectorizeOpPrecondition(Operation *op,
                                       ArrayRef<int64_t> inputVectorSizes,
                                       bool vectorizeNDExtract) {
   return TypeSwitch<Operation *, LogicalResult>(op)
       .Case<linalg::LinalgOp>([&](auto linalgOp) {

mlir/test/Dialect/Linalg/vectorization-masked.mlir

 
 // -----
 
 // CHECK-LABEL: func @test_masked_vectorize_pad
 func.func @test_masked_vectorize_pad(
   %0 : tensor<?x?xf32>, %h0 : index, %h1 : index)
     -> tensor<2x4xf32>
 {
-  //  CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
-  //  CHECK-DAG: %[[c0_2:.*]] = arith.constant 0 : index
   //  CHECK-DAG: %[[c42:.*]] = arith.constant 4.243000e+01 : f32
+  //  CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
   //  CHECK-DAG: %[[empty:.*]] = tensor.empty() : tensor<2x4xf32>
   //      CHECK: %[[d0:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>
   //      CHECK: %[[d1:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>
   //      CHECK: %[[mask:.*]] = vector.create_mask %[[d0]], %[[d1]] : vector<2x4xi1>
+  //  CHECK-DAG: %[[c0_2:.*]] = arith.constant 0 : index
   //      CHECK: %[[masked_read:.*]] = vector.mask %[[mask]] {
   // CHECK-SAME:   vector.transfer_read %{{.*}}[%[[c0_2]], %[[c0_2]]], %[[c42]]
   // CHECK-SAME:   {in_bounds = [true, true]} : tensor<?x?xf32>, vector<2x4xf32>
   // CHECK-SAME: } : vector<2x4xi1> -> vector<2x4xf32>
   //      CHECK: vector.transfer_write %[[masked_read]], %[[empty]][%[[c0_2]], %[[c0_2]]]
   // CHECK-SAME:   {in_bounds = [true, true]} : vector<2x4xf32>, tensor<2x4xf32>
   %cst = arith.constant 42.43 : f32
   %c0 = arith.constant 0 : index
   %1 = tensor.pad %0 low[0, %c0] high[%h0, %h1]  {
     ^bb0(%hh1: index, %hh2: index):
       tensor.yield %cst : f32
     } : tensor<?x?xf32> to tensor<2x4xf32>
   return %1: tensor<2x4xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["tensor.pad"]} in %arg1
     : (!transform.any_op) -> !transform.any_op
   transform.structured.masked_vectorize %0 vector_sizes [2, 4] : !transform.any_op
 }
 
 // -----
 
+// CHECK-LABEL: func @test_masked_vectorize_dynamic_pad
+func.func @test_masked_vectorize_dynamic_pad(
+  %0 : tensor<?x?xf32>, %h0 : index, %h1 : index)
+    -> tensor<?x?xf32>
+{
+  //  CHECK-DAG: %[[c42:.*]] = arith.constant 4.243000e+01 : f32
+  //  CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
+  //      CHECK: %[[empty:.*]] = tensor.empty({{.+}}) : tensor<?x?xf32>
+  //      CHECK: %[[d0:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>
+  //      CHECK: %[[d1:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>
+  //      CHECK: %[[mask:.*]] = vector.create_mask %[[d0]], %[[d1]] : vector<2x4xi1>
+  //  CHECK-DAG: %[[c0_2:.*]] = arith.constant 0 : index
+  //      CHECK: %[[masked_read:.*]] = vector.mask %[[mask]] {
+  // CHECK-SAME:   vector.transfer_read %{{.*}}[%[[c0_2]], %[[c0_2]]], %[[c42]]
+  // CHECK-SAME:   {in_bounds = [true, true]} : tensor<?x?xf32>, vector<2x4xf32>
+  // CHECK-SAME: } : vector<2x4xi1> -> vector<2x4xf32>
+  //      CHECK: %[[masked_write:.*]] = vector.mask %[[mask]] {
+  // CHECK-SAME: vector.transfer_write %[[masked_read]], %[[empty]][%[[c0_2]], %[[c0_2]]]
+  // CHECK-SAME:   {in_bounds = [true, true]} : vector<2x4xf32>, tensor<?x?xf32>
+  //      CHECK: return %[[masked_write]] : tensor<?x?xf32>
+  %cst = arith.constant 42.43 : f32
+  %c0 = arith.constant 0 : index
+  %1 = tensor.pad %0 low[0, %c0] high[%h0, %h1]  {
+    ^bb0(%hh1: index, %hh2: index):
+      tensor.yield %cst : f32
+    } : tensor<?x?xf32> to tensor<?x?xf32>
+  return %1: tensor<?x?xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["tensor.pad"]} in %arg1
+    : (!transform.any_op) -> !transform.any_op
+  transform.structured.masked_vectorize %0 vector_sizes [2, 4] : !transform.any_op
+}
+
+// -----
+
 func.func @vectorize_dynamic_matmul(%A: memref<?x?xf32>, %B: memref<?x?xf32>, %C: memref<?x?xf32>) {
   linalg.matmul ins(%A, %B: memref<?x?xf32>, memref<?x?xf32>)
             outs(%C: memref<?x?xf32>)
   return
 }
 
 // CHECK-LABEL:   func.func @vectorize_dynamic_matmul(
 // CHECK-SAME:      %[[VAL_0:.*]]: memref<?x?xf32>, %[[VAL_1:.*]]: memref<?x?xf32>, %[[VAL_2:.*]]: memref<?x?xf32>) {