[Mlir-commits] [mlir] [MLIR] Allowing unsupported conv2d op to fail gracefully from vectorization (PR #130181)

[Andrzej Warzyński]

================
@@ -1990,8 +1996,21 @@ static LogicalResult vectorizeLinalgOpPrecondition(
   // 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>(linalgOp.getOperation()))
+  if (isa<ConvolutionOpInterface>(linalgOp.getOperation())) {
+    // Create a dummy rewriter first, a rewriter is not required for
+    // validation
+    IRRewriter dummyBuilder(linalgOp.getContext());
+    // Check if we can successfully construct a 1d convolution generator.
+    // For example, if it is 2d+ convolution, return failure because we don't
+    // support it. To use this pass on a 2d+ convolution, it should have already
+    // been decomposed to 1d convolution via
+    // DecomposeConvolutionToLowerDimOpsPass.
----------------
banach-space wrote:

> Let me use this thread to discuss it. This diff code block is the high-level pre condition check around linalg op. Please let me know if you are referring to a different location.

Similarly to Diego, I am suggesting to move the high level logic to `vectorizeOpPrecondition`. Also, like other "pre-condition" hooks, it should not require a rewriter.

> Could you elaborate? Are you referring to explicitly invoke `inferConvolutionDims()`? Then to make sure this is a regular 2d convolution, I'd check for a combination of:
> 
> * outputImage.size() == 2
> * batch.size() == 1
> * outputChannel.size() == 1
> 
> Reject if all of those satisfy. I have no problem to implement this, but just want to make sure we are on the same page.

In my naivety, I was hoping that checking e.g. the rank of the filter or the input would be sufficient. But clearly not - the input for non-channelled conv would be 1D, but for a channeled one would be 2D. So on and so forth. IMO, you can just create something like this:

```cpp
if (!isa<conv_type1, conv_type2, ...>(conv))
   return failure();
```

This will be a bit verbose, but there's just too many convs and whatever we try will be ... verbose 🤷🏻 

> 
> Taking a step back, I don't have a lot of context about the history of the vectorization code around convolution. Since this PR is not intending to do a massive re-write, I'm attempting to be coherent with the existing code as much as possible.
> 

+1 to being coherent, thanks! 

I was actually going to ask - do you have any plans regarding this code beyond this PR? 

> One thing I've noticed and @hanhanW who righteously pointed out is that we can fail to build a `Conv1DGenerator`, and still allow a function (like how `vectorizeConvolution()` construct and uses the `Conv1DGenerator`) invoked on its member vectorization functions, which I find to be quite confusing. (If I'm to implement this from scratch, I'll probably use singleton + initialize compared to the approach (constructor + `valid` member variable). This way, a developer is required to invoke the initialize method and check validity of the class before invoking anything on it.)

You should be able to simply add:
```cpp
assert(isValid() && "Conv1DGenerator failed")
```

>From what I can tell, that wouldn't break any tests and will make "validity" a strong pre-requisite.

> 
> With this context, I find the most defensive approach is the one used from this PR right now:
> 
> * With future implementation to be added and more flavor of convolution supported, it is very likely that the precondition check on vectorize convolution grow out of sync (and this PR is a perfect example)

There's been no new implementations in > 2 yrs. From what I can tell, we can safely assume that this will remain the case for the foreseeable future. So, I wouldn't worry about this.

> * Now instead of maintain a separate function that does a subset of the constructor logic, why not re-use it and ensure we do the validity check? This looks reasonable as the constructor is (if not better, at least) not more expensive than having to infer the convolution dimensions.

That sounds good in theory, but in practice it means that we need an IR writer for the validation. "Validation"/"pre-conditioning" should not require a rewriter.

> 
> With above reasoning added up, it just looks to me to be a better solution compared with inferring the convolution dimensions and reject a few corner cases (which can easily grow out-of-sync later).

How about my suggestion with `isa`?

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