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

[Zhuoran Yin]

================
@@ -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.
----------------
jerryyin wrote:

Apologize, this is an iree implementation details and referenced pass is also from IREE. I shouldn't include such pass in upstream. Will remove.

> I am proposing to re-use the high-level logic that's already available.

I'm using this thread to discuss this point this the diff in this block is the high-level pre condition check around linalg op. Please let me know if you are referring to a different location.

> From what I can tell, it would be totally fine to check the dims of conv ops very early on.

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.

------

Taking a step back, I don't have a lot of context of 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. 

One thing I've noticed and @hanhanW who righteously pointed out is that we can fail to build a `Conv1DGenerator`, and still allow a member function (like how `vectorizeConvolution()` construct and uses the `Conv1DGenerator`) invoked to vectorize it, 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).  This way, a developer is required to invoke the initialize method and check validity of the class before invoking anything on it.)

With this context, I find the most defensive approach is the one used from this PR: 
 - With future implemented 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)
 - 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 fairly cheap, at least not worse than having to infer the convolution dimensions.

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 will grow out-of-sync later).

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