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

[mravishankar]

Summary

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.

[dcaballe]

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

571

nit: regionOps -> regionOp

572

Don't we vectorize some tensor extract ops already?

[bondhugula]

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

571

You don't need getOperations().

571

Nit: The name regionOps isn't appropriate here.

[mravishankar]

Address comments.

[mravishankar]

Rebase

[mravishankar]

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

571

Used payloadOp instead.

572

I think not all tensor.extract operations are vectorized. This might need to be made tighter if needed. I'd just wait till all tensor.extract ops are vectorized to drop this here.

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 //===----------------------------------------------------------------------===//
 // Preconditions that ensure the corresponding transformation succeeds and can
 // be applied as a rewrite pattern.
 //===----------------------------------------------------------------------===//
 /// Promote memref.subviews feeding linalg-on-buffers operations.
 LogicalResult promoteSubviewsPrecondition(Operation *op,
                                           LinalgPromotionOptions options);
 
+/// Return success if the operation can be vectorized.
+LogicalResult vectorizeLinalgOpPrecondition(LinalgOp linalgOp);
+
 //===----------------------------------------------------------------------===//
 // Transformations exposed as rewrite patterns.
 //===----------------------------------------------------------------------===//
 // Marker used as attribute name in generated Linalg rewriting transformations.
 struct LinalgTransforms {
   static const StringLiteral kLinalgTransformMarker;
 };
 

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

if (llvm::any_of(innerOp.getOperandTypes(), [](Type type) {
       return failure();
     }
     if (llvm::any_of(innerOp.getResultTypes(), [](Type type) {
           return !VectorType::isValidElementType(type);
         })) {
       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...
   if (isa<ConvolutionOpInterface>(op.getOperation()))
     return success();
   // TODO: the common vector shape is equal to the static loop sizes only when
   // all indexing maps are projected permutations. For convs and stencils the
   // logic will need to evolve.
   if (!allIndexingsAreProjectedPermutation(op)) {
     LDBG("precondition failed: not projected permutations");
     return failure();
   }
   if (failed(reductionPreconditions(op))) {
     LDBG("precondition failed: reduction preconditions");
     return failure();
   }
   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");
     return failure();
   }
   return vectorizeStaticLinalgOpPrecondition(linalgOp);
 }