[GlobalOptimization] Add opt-in flag for Conv2D to Img2Col transformation

Add --iree-global-opt-enable-conv2d-to-img2col flag to enable the
transformation. The pass is disabled by default and runs after
LinalgQuantizedConvToConvPass in the global optimization pipeline.

Signed-off-by: Phoebe Chen <phoebe.chen@sifive.com>

Author: phoebesv

compiler/src/iree/compiler/GlobalOptimization/Passes.cpp

@@ -118,6 +118,8 @@ void buildGlobalOptimizationPassPipeline(
       .addPass(IREE::Util::createOptimizeIntArithmeticPass)
       .addPass(createLinalgQuantizedConvToConvPass)
       .addPass(createLinalgQuantizedMatmulToMatmulPass)
+      .addPredicatedPass(transformOptions.enableConv2DToImg2Col,
+                         createConvertConv2DToImg2ColPass)
       .addPass(IREE::Flow::createCanonicalizePass)
       .addPass(createRemoveZeroExtentTensorsPass)
       .addPass(createDetachElementwiseFromNamedOpsPass)

compiler/src/iree/compiler/Pipelines/Options.cpp

@@ -302,6 +302,11 @@ void GlobalOptimizationOptions::bindOptions(OptionsBinder &binder) {
       llvm::cl::desc("Convert named matmul ops to linalg generic ops during "
                      "global optimization to enable better fusion."),
       llvm::cl::cat(category));
+  binder.opt<bool>(
+      "iree-global-opt-enable-conv2d-to-img2col", enableConv2DToImg2Col,
+      llvm::cl::desc("Enables conversion of Conv2D operations to img2col + "
+                     "matmul form to leverage optimized matmul implementations."),
+      llvm::cl::cat(category));
 }
 
 void SchedulingOptions::bindOptions(OptionsBinder &binder) {

compiler/src/iree/compiler/Pipelines/Options.h

@@ -187,6 +187,9 @@ struct GlobalOptimizationOptions {
   // Converts linalg named matmul ops to linalg generic ops.
   bool generalizeMatmul = false;
 
+  // Enables conversion of Conv2D operations to img2col + matmul form.
+  bool enableConv2DToImg2Col = false;
+
   void bindOptions(OptionsBinder &binder);
   using FromFlags = OptionsFromFlags<GlobalOptimizationOptions>;
 };

compiler/src/iree/compiler/Pipelines/Pipelines.cpp

@@ -212,6 +212,8 @@ void buildIREEPrecompileTransformPassPipeline(
       globalOptimizationOptions.stripAssertions;
   globalTransformOptions.generalizeMatmul =
       globalOptimizationOptions.generalizeMatmul;
+  globalTransformOptions.enableConv2DToImg2Col =
+      globalOptimizationOptions.enableConv2DToImg2Col;
   globalTransformOptions.constExprHoisting = pipelineOptions.constExprHoisting;
   globalTransformOptions.constExprMaxSizeIncreaseThreshold =
       pipelineOptions.constExprMaxSizeIncreaseThreshold;