Forward declare Conv1DGenerator for validaty

Author: jerryyin

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

@@ -52,6 +52,12 @@ using namespace mlir::linalg;
 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
 #define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")
 
+// Forward declaration of Conv1DGenerator and its validator
+namespace {
+struct Conv1DGenerator;
+bool validateConv1DGenerator(RewriterBase &rewriter, LinalgOp linalgOp);
+} // namespace
+
 /// Try to vectorize `convOp` as a convolution.
 static FailureOr<Operation *>
 vectorizeConvolution(RewriterBase &rewriter, LinalgOp convOp,
@@ -1991,14 +1997,17 @@ static LogicalResult vectorizeLinalgOpPrecondition(
   // features. But we will still need stride/dilation attributes that will be
   // annoying to reverse-engineer...
   if (isa<ConvolutionOpInterface>(linalgOp.getOperation())) {
-    // Check if it is 2d+ convolution. If it is, return failure because we don't
+    // 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.
-    if (linalgOp.getNumParallelLoops() >= 4) {
-      LDBG("precondition failed: Regular 2d+ convolutions not supported.\n");
+    if (!validateConv1DGenerator(dummyBuilder, linalgOp))
       return failure();
-    }
+
     return success();
   }
 
@@ -3197,6 +3206,8 @@ struct Conv1DGenerator
     valid = true;
   }
 
+  bool isValid() { return valid; }
+
   /// Generate a vector implementation for:
   /// ```
   ///   Op def: (     w,     kw  )
@@ -3907,6 +3918,13 @@ struct Conv1DGenerator
     }
   }
 };
+
+// Helper function to construct Conv1DGenerator
+bool validateConv1DGenerator(RewriterBase &rewriter, LinalgOp linalgOp) {
+  Conv1DGenerator conv1dGen(rewriter, linalgOp);
+  return conv1dGen.isValid();
+}
+
 } // namespace
 
 /// Helper function to vectorize a LinalgOp with convolution semantics.