[Mlir-commits] [mlir] b858309 - [mlir] Only attempt to vectorize conv if conv.

[Jacques Pienaar]

Author: Jacques Pienaar
Date: 2023-10-18T20:34:39-07:00
New Revision: b858309ddc977d5e70de54f3fa3888915b5fbc0c

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

LOG: [mlir] Only attempt to vectorize conv if conv.

Avoids hitting assertions due to unsupported convolution patterns.

See https://github.com/openxla/iree/issues/15207#issuecomment-1767650797

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index d60a1a1cda849c7..b427af33e3c4400 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -1693,11 +1693,16 @@ LogicalResult mlir::linalg::vectorize(RewriterBase &rewriter, Operation *op,
           .Case<linalg::LinalgOp>([&](auto linalgOp) {
             // TODO: isaConvolutionOpInterface that can also infer from generic
             // features. Will require stride/dilation attributes inference.
-            FailureOr<Operation *> convOr =
-                vectorizeConvolution(rewriter, linalgOp);
-            if (succeeded(convOr)) {
-              llvm::append_range(results, (*convOr)->getResults());
-              return success();
+            if (isa<ConvolutionOpInterface>(linalgOp.getOperation())) {
+              FailureOr<Operation *> convOr =
+                  vectorizeConvolution(rewriter, linalgOp);
+              if (succeeded(convOr)) {
+                llvm::append_range(results, (*convOr)->getResults());
+                return success();
+              }
+
+              LDBG("Unsupported convolution can't be vectorized.\n");
+              return failure();
             }
 
             LDBG("Vectorize generic by broadcasting to the canonical vector "