[Mlir-commits] [mlir] bf5f63e - [memref][Transform][NFC] Improve the doc for masked_vectorize

[Quentin Colombet]

Author: Quentin Colombet
Date: 2023-01-27T11:20:05+01:00
New Revision: bf5f63e59fd729efd6dd69318b365293f7ce385d

URL: https://github.com/llvm/llvm-project/commit/bf5f63e59fd729efd6dd69318b365293f7ce385d
DIFF: https://github.com/llvm/llvm-project/commit/bf5f63e59fd729efd6dd69318b365293f7ce385d.diff

LOG: [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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
    mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index bd2fe3b4cf84a..043dc778b00fb 100644
--- 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 b356de9bbdf86..3173a447b5756 100644
--- 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,