diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp @@ -1445,10 +1445,23 @@ } ArrayRef resultTensorShape = padOp.getResultType().getShape(); - if (!(resultTensorShape == inputVectorSizes)) { - LDBG("result tensor shape must match input vector sizes: " << padOp - << "\n"); - // return failure(); + 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) {