[memref][Transform][NFC] Improve the doc for masked_vectorize

The `transform.structured.masked_vectorize` operator assumes that the
iteration space of the given linalg op is smaller than the given vector
sizes.
Explicitly states this requirement in the description of the operation.

Also fix the related comment and assert message in the vectorization code.
The wording was flipped.

NFC

Differential Revision: https://reviews.llvm.org/D142628

diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index bd2fe3b4cf84ab3e6dec3a56ce8e0a5ce2206b25..043dc778b00fb645c21134b144341c2402205b3b 100644 (file)
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
@@ -1507,6 +1507,12 @@ def MaskedVectorizeOp : Op<Transform_Dialect, "structured.masked_vectorize",
     SSA values, the handle must be mapped to exactly one payload op with
     exactly one index-typed result.
 
+    Note: The input vector sizes must be bigger than or equal to their
+    counterpart iteration space sizes.
+
+    Typically this operator should be applied to linalg operations that have
+    already be tiled to the appropriate sizes.
+
     #### Return modes:
 
     This operation produces a definite failure if the dynamic vector sizes (SSA

diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index b356de9bbdf86aa0f2c90396d08a74fdb3fe55e7..3173a447b575659bf9d210b4d49dc5ec80822ee4 100644 (file)
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -980,16 +980,17 @@ mlir::linalg::vectorizeLinalgOpPrecondition(LinalgOp linalgOp,
            "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),
-               [](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 smaller or equal than iteration space "
-           "static sizes");
+    assert(
+        llvm::all_of(
+            llvm::zip(linalgOp.getStaticLoopRanges(), 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 "
+        "static sizes");
   }
 
   // TODO: Masking is only supported for dynamic shapes so input vector sizes
@@ -1066,8 +1067,8 @@ static void convertAffineApply(RewriterBase &rewriter, LinalgOp linalgOp) {
 
 /// Emit a suitable vector form for a Linalg op. If provided, `inputVectorSizes`
 /// are used to vectorize this operation. `inputVectorSizes` must match the rank
-/// of the iteration space of the operation and the sizes must be smaller or
-/// equal than their counterpart interation space sizes, if static.
+/// of the iteration space of the operation and the input vector sizes must be
+/// greater than or equal to their counterpart iteration space sizes, if static.
 /// `inputVectorShapes` also allows the vectorization of operations with dynamic
 /// shapes.
 LogicalResult mlir::linalg::vectorize(RewriterBase &rewriter, LinalgOp linalgOp,