[Mlir-commits] [mlir] a0ef8af - [mlir][Linalg] Expose vectorization precondition check as a utility function.

[Mahesh Ravishankar]

Author: Mahesh Ravishankar
Date: 2022-10-05T17:24:52Z
New Revision: a0ef8af8d5d04c8a7b18cd1ce4c2a728d2f480b2

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

LOG: [mlir][Linalg] Expose vectorization precondition check as a utility function.

This patch exposes the method to check if an op can be vectorized or
not for downstream uses. Also adds a check to mark elementwise operations
that have non-vectorizable ops (like `tensor.extract`) as non vectorizable.

Reviewed By: nicolasvasilache, dcaballe, ThomasRaoux

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
    mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 727c8d48491c8..6d044ed489b33 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -344,6 +344,9 @@ FailureOr<LinalgLoops> linalgOpToAffineLoops(PatternRewriter &rewriter,
 LogicalResult promoteSubviewsPrecondition(Operation *op,
                                           LinalgPromotionOptions options);
 
+/// Return success if the operation can be vectorized.
+LogicalResult vectorizeLinalgOpPrecondition(LinalgOp linalgOp);
+
 //===----------------------------------------------------------------------===//
 // Transformations exposed as rewrite patterns.
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index cc98013971cdb..5623a16fb2613 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -566,8 +566,16 @@ static LogicalResult vectorizeStaticLinalgOpPrecondition(linalg::LinalgOp op) {
       return failure();
     }
   }
-  if (isElementwise(op))
+  if (isElementwise(op)) {
+    // Some operations in the body cannot be vectorized.
+    for (Operation &payloadOp : *op.getBlock()) {
+      if (isa<tensor::ExtractOp>(payloadOp)) {
+        LDBG("precondition failed: `tensor.extract` not vectorizable");
+        return failure();
+      }
+    }
     return success();
+  }
   // TODO: isaConvolutionOpInterface that can also infer from generic features.
   // But we will still need stride/dilation attributes that will be annoying to
   // reverse-engineer...
@@ -587,7 +595,7 @@ static LogicalResult vectorizeStaticLinalgOpPrecondition(linalg::LinalgOp op) {
   return success();
 }
 
-static LogicalResult vectorizeLinalgOpPrecondition(LinalgOp linalgOp) {
+LogicalResult mlir::linalg::vectorizeLinalgOpPrecondition(LinalgOp linalgOp) {
   // All types must be static shape to go to vector.
   if (linalgOp.hasDynamicShape()) {
     LDBG("precondition failed: dynamic shape");