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

[Andrzej Warzyński]

banach-space wrote:

> It is confusing because we have `vectorizeOpPrecondition` method; the real checks happen there.
> 
> https://github.com/llvm/llvm-project/blob/ce72c76e9bd0bf74af614cae6a9f85cfd4720e95/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp#L2091-L2115

Yes, there are effectively two levels of verification (note that this PR doesn't change that). Here's pseudo-code to explain this:

```cpp
// Verification 1: At the level of transform.structured.vectorize (very crude, initial filtering)
LogicalResult process_vectorize_td(Operation *target) {
  if (!isVectorizable)
    return rewriter.notifyMatchFailure(op, "Unsupported Op, cannot vectorize");

  return vectorize(target);
}

// Verification 2: At the level of linalg::vectorize (very detailed)
LogicalResult vectorize(Operation *op) {
    if (failed(vectorizeOpPrecondition(op) {
      LDBG("Vectorization pre-conditions failed\n");
      return failure();
    }
    // ...
}
```

Basically:
* first, before doing anything expensive (and entering `vectorize()`), we check that the candidate Op is supported (*),
* second, for every Op that's supported, we run more involved checks inside `vectorize()`.

I can re-wire things, but would rather start with this small change :) I have a few other changes in flight that I'd like to prioritise (related to vectorising `tensor.pad`/`tensor.pack`).

(*) You are right that even some Linalg Ops don't qualify for vectorization (at least today). The idea behind `isVectorizable` (I should probably rename it) is to reject Ops that are known not to be supported at all, i.e. Ops outside the Linalg Dialect and selected Ops from Tensor. And for the supported Ops, as you pointed out, we still need to run the pre-condition checks and reject some of them.

Hope I didn't miss anything 😅 

https://github.com/llvm/llvm-project/pull/110708