[Mlir-commits] [mlir] [mlir][linalg] Add a new helper hook - `isVectorizable` (PR #110708)

[Andrzej Warzyński]

https://github.com/banach-space created https://github.com/llvm/llvm-project/pull/110708

The newly added hook simply returns `false` for Ops for which there's no
"vectorization logic" in the Linalg Vectorizer (i.e. the `vectorize`
method). It's added so that the following two TD ops expose identical
level of functionality (that's not the case ATM):

  * `transform.structured.vectorize_children_and_apply_patterns`
  * `transform.structured.vectorize`


>From 2a0be5e7c3763f680f40293686c350d04c4174ee Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Tue, 1 Oct 2024 18:33:29 +0100
Subject: [PATCH] [mlir][linalg] Add a new helper hook - `isVectorizable`

The newly added hook simply returns `false` for Ops for which there's no
"vectorization logic" in the Linalg Vectorizer (i.e. the `vectorize`
method). It's added so that the following two TD ops expose identical
level of functionality (that's not the case ATM):

  * `transform.structured.vectorize_children_and_apply_patterns`
  * `transform.structured.vectorize`
---
 .../mlir/Dialect/Linalg/Transforms/Transforms.h     |  7 +++++++
 .../Linalg/TransformOps/LinalgTransformOps.cpp      | 13 ++++++-------
 .../lib/Dialect/Linalg/Transforms/Vectorization.cpp |  5 +++++
 3 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 48e657cca96e39..facda13eeaa7ed 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -762,6 +762,13 @@ LogicalResult copyToGPUPrivateMemory(OpBuilder &b, Value src, Value dst);
 /// memory is freed when going outside of the scope.
 LogicalResult deallocateGPUPrivateMemory(OpBuilder &, Value /*buffer*/);
 
+/// Check if this Op is vectorizable. All Linalg Op are vectorizable, as well
+/// as selected Tensor Ops. Note that this is merely a high level check and
+/// that the vectorizer also requires various additional pre-conditions to be
+/// met for it to work. These are only checked for Ops that are supported,
+/// other Ops should be rejected early.
+bool isVectorizable(Operation *);
+
 /// Emit a suitable vector form for an operation. If provided,
 /// `inputVectorSizes` are used to vectorize this operation. `inputVectorSizes`
 /// must match the rank of the iteration space of the operation and the sizes
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 46c8510f4ed514..8632e7db76353c 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -3411,11 +3411,11 @@ struct VectorizationPattern : public RewritePattern {
         flatten1DDepthwiseConv(flattenConv) {}
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
-    LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
-    if (!linalgOp)
-      return rewriter.notifyMatchFailure(op, "expected Linalg Op");
-    return vectorize(rewriter, linalgOp, /*inputVectorSizes=*/{},
-                     /*scalableVecDims=*/{}, vectorizeNDExtract,
+    if (!linalg::isVectorizable(op))
+      return rewriter.notifyMatchFailure(op,
+                                         "Unsupported Op, cannot vectorize");
+    return vectorize(rewriter, op, /*inputVectorSizes=*/{},
+                     /*inputScalableVecDims=*/{}, vectorizeNDExtract,
                      flatten1DDepthwiseConv);
   }
 
@@ -3496,8 +3496,7 @@ DiagnosedSilenceableFailure transform::VectorizeOp::apply(
 
   // TODO: Check that the correct number of vectorSizes was provided.
   for (Operation *target : targets) {
-    if (!isa<linalg::LinalgOp, tensor::PadOp, tensor::PackOp, tensor::UnPackOp>(
-            target)) {
+    if (!linalg::isVectorizable(target)) {
       return mlir::emitSilenceableFailure(target->getLoc())
              << "Unsupported Op, cannot vectorize";
     }
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index ca85f4b9b9c156..d873ca41ee86f8 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -2129,6 +2129,11 @@ static void convertAffineApply(RewriterBase &rewriter, LinalgOp linalgOp) {
   }
 }
 
+bool mlir::linalg::isVectorizable(Operation *op) {
+  return isa<linalg::LinalgOp, tensor::PadOp, tensor::PackOp, tensor::UnPackOp>(
+      op);
+}
+
 /// Emit a suitable vector form for an operation. If provided,
 /// `inputVectorSizes` are used to vectorize this operation.
 /// `inputVectorSizes` must match the rank of the iteration space of the