[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

1326

Could you add a test to exercise this?

1455

[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.

1460–1461

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?).

1468

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

[dcaballe]

Awesome! Just a few comments. LGTM overall

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

1326

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?

1454–1473

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

1326

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.

1455

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.

1460–1461

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

1460–1461

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,
   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 destType = RankedTensorType::get(inputVectorSizes, padValue.getType());
   auto emptyOp =
-      rewriter.create<tensor::EmptyOp>(loc, padOp.getResultType(),
+      rewriter.create<tensor::EmptyOp>(loc, destType,
                                        /*dynamicSizes=*/ValueRange{});
   SmallVector<OpFoldResult> mixedSourceDims =
       getMixedDimensions(rewriter, loc, padOp.getSource());
   Value mask =
       rewriter.create<vector::CreateMaskOp>(loc, maskType, mixedSourceDims);
   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>(
       loc,
       /*vector=*/maskedOp->getResult(0),
       /*source=*/emptyOp,
       /*indices=*/SmallVector<Value>(rank, zero),
       /*inBounds=*/SmallVector<bool>(rank, true));
-  newResults.push_back(transferWriteOp.getResult());
+  Value result = transferWriteOp.getResult();
+  if (destType != padOp.getResultType()) {
+    result =
+        rewriter.create<tensor::CastOp>(loc, padOp.getResultType(), result);
+  }
+  newResults.push_back(result);
   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");
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();
 }
 
 LogicalResult
 mlir::linalg::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.
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
+  if (inputVectorSizes.size() != resultTensorShape.size()) {
+    LDBG("Input vector sizes don't match the rank of pad op : " << padOp
                                                                 << "\n");
     return failure();
   }
 
+  assert(!ShapedType::isDynamicShape(inputVectorSizes) &&
+         "Input vector sizes can't have dynamic dimensions");
+
+  for (auto [dimSize, vecSize] :
+       llvm::zip_equal(resultTensorShape, inputVectorSizes)) {
+    if (ShapedType::isDynamic(dimSize)) continue;
+    if (vecSize < dimSize) {
+      LDBG("input vector sizes must be greater than or equal to result sizes: "
+           << padOp << "\n");
+      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();

mlir/test/Dialect/Linalg/vectorization.mlir

 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["tensor.pad"]} in %arg1
     : (!pdl.operation) -> !pdl.operation
   transform.structured.masked_vectorize %0 vector_sizes [2, 4]
 }
 
 // -----
 
+// 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: %[[c0:.*]] = arith.constant 0 : index
+  //  CHECK-DAG: %[[c0_2:.*]] = arith.constant 0 : index
+  //  CHECK-DAG: %[[c42:.*]] = arith.constant 4.243000e+01 : f32
+  //  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: %[[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<?x?xf32>
+  return %1: tensor<?x?xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["tensor.pad"]} in %arg1
+    : (!pdl.operation) -> !pdl.operation
+  transform.structured.masked_vectorize %0 vector_sizes [2, 4]
+}
+
+// -----
+
 // CHECK-LABEL: func @test_masked_pad_static_dynamic
 func.func @test_masked_pad_static_dynamic(%arg0: tensor<1x2x2x?xf32>, %low: index, %high: index,
                   %pad_value: f32) -> tensor<6x?x?x?xf32> {
   // CHECK: tensor.pad
   %0 = tensor.pad %arg0 low[2, %low, 3, 3] high[3, 3, %high, 2] {
     ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
       tensor.yield %pad_value : f32
     } : tensor<1x2x2x?xf32> to tensor<6x?x?x?xf32>