[LLVMCPU] Add an additional level of tiling

Add an addition level of tiling aka vector parallel to the CPU default
pipeline. For some of the linalg op, that is not specialized through any
pipeline may hit bufferization issue, if passed through the default
pipeline. Adding an extra level of tiling takes care of such cases.

compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp

@@ -2506,6 +2506,13 @@ static LogicalResult setRootConfig(mlir::FunctionOpInterface entryPointFn,
   SmallVector<int64_t> distTileSizes =
       getDefaultDistributedLevelTileSizes(op, DistributionHeuristicConfig{});
   TileSizesListType tileSizes = {distTileSizes};
+
+  // Add an extra level of tiling.
+  if (auto linalgOp = dyn_cast<linalg::LinalgOp>(*op)) {
+    SmallVector<int64_t> vecTileSizes = distTileSizes;
+    limitVectorTileSizes(linalgOp, vecTileSizes);
+    tileSizes.push_back(vecTileSizes);
+  }
   return setOpConfigAndEntryPointFnTranslation(
       entryPointFn, op, tileSizes, DispatchLoweringPassPipeline::CPUDefault);
 }

compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPULowerExecutableTarget.cpp

@@ -64,29 +64,20 @@ class LLVMCPULowerExecutableTargetPass
 };
 } // namespace
 
-// TODO(dcaballe): We temporarily need this utility to retrieve a valid
-// lowering config. We should be able to remove this once we have a lowering
-// config attribute per op.
-static FailureOr<LoweringConfigAttr>
-getRootLoweringConfig(FunctionOpInterface funcOp) {
+static FailureOr<TilingConfig>
+getTilingConfigForPipeline(FunctionOpInterface funcOp) {
   SmallVector<Operation *> computeOps = getComputeOps(funcOp);
-  // Check for self first.
   FailureOr<Operation *> rootOp = getRootOperation(computeOps);
+  if (failed(rootOp) || !rootOp.value()) {
+    return failure();
+  }
   auto rootLoweringConfig =
       iree_compiler::getLoweringConfig<IREE::Codegen::LoweringConfigAttr>(
           rootOp.value());
-  if (rootLoweringConfig) {
-    return rootLoweringConfig;
+  if (!rootLoweringConfig) {
+    return failure();
   }
-
-  return failure();
-}
-
-static TilingConfig getTilingConfigForPipeline(FunctionOpInterface funcOp) {
-  auto maybeLoweringConfig = getRootLoweringConfig(funcOp);
-  assert(succeeded(maybeLoweringConfig) &&
-         "Pipeline requires a lowering config");
-  return TilingConfig(*maybeLoweringConfig);
+  return TilingConfig(rootLoweringConfig);
 }
 
 void LLVMCPULowerExecutableTargetPass::runOnOperation() {
@@ -122,42 +113,76 @@ void LLVMCPULowerExecutableTargetPass::runOnOperation() {
   // No pipleline specified, nothing to do.
   case IREE::Codegen::DispatchLoweringPassPipeline::None:
     return;
-  case IREE::Codegen::DispatchLoweringPassPipeline::CPUDefault:
-    addCPUDefaultPassPipeline(pipeline);
+  case IREE::Codegen::DispatchLoweringPassPipeline::CPUDefault: {
+    auto maybeTilingConfig = getTilingConfigForPipeline(funcOp);
+    addCPUDefaultPassPipeline(pipeline, maybeTilingConfig);
     break;
+  }
   case IREE::Codegen::DispatchLoweringPassPipeline::
       CPUBufferOpsTileAndVectorize: {
-    TilingConfig tilingConfig = getTilingConfigForPipeline(funcOp);
-    addCPUBufferOpsTileAndVectorizePipeline(pipeline, tilingConfig,
+    auto maybeTilingConfig = getTilingConfigForPipeline(funcOp);
+    if (failed(maybeTilingConfig)) {
+      funcOp.emitOpError("Tiling Config is necessary for "
+                         "CPUBufferOpsTileAndVectorize pipeline.");
+      return signalPassFailure();
+    }
+    addCPUBufferOpsTileAndVectorizePipeline(pipeline, *maybeTilingConfig,
                                             pipelineOpts);
     break;
   }
   case IREE::Codegen::DispatchLoweringPassPipeline::CPUDoubleTilingExpert: {
-    TilingConfig tilingConfig = getTilingConfigForPipeline(funcOp);
-    addMultiTilingExpertPassPipeline(pipeline, tilingConfig, pipelineOpts);
+    auto maybeTilingConfig = getTilingConfigForPipeline(funcOp);
+    if (failed(maybeTilingConfig)) {
+      funcOp.emitOpError(
+          "Tiling Config is necessary for CPUDoubleTilingExpert pipeline.");
+      return signalPassFailure();
+    }
+    addMultiTilingExpertPassPipeline(pipeline, *maybeTilingConfig,
+                                     pipelineOpts);
     break;
   }
   case IREE::Codegen::DispatchLoweringPassPipeline::
       CPUConvTileAndDecomposeExpert: {
-    TilingConfig tilingConfig = getTilingConfigForPipeline(funcOp);
-    addConvTileAndDecomposeExpertPassPipeline(pipeline, tilingConfig,
+    auto maybeTilingConfig = getTilingConfigForPipeline(funcOp);
+    if (failed(maybeTilingConfig)) {
+      funcOp.emitOpError("Tiling Config is necessary for "
+                         "CPUConvTileAndDecomposeExpert pipeline.");
+      return signalPassFailure();
+    }
+    addConvTileAndDecomposeExpertPassPipeline(pipeline, *maybeTilingConfig,
                                               pipelineOpts);
     break;
   }
   case IREE::Codegen::DispatchLoweringPassPipeline::Mmt4dTilingExpert: {
-    TilingConfig tilingConfig = getTilingConfigForPipeline(funcOp);
-    addMmt4dTilingExpertPassPipeline(pipeline, tilingConfig, pipelineOpts);
+    auto maybeTilingConfig = getTilingConfigForPipeline(funcOp);
+    if (failed(maybeTilingConfig)) {
+      funcOp.emitOpError(
+          "Tiling Config is necessary for Mmt4dTilingExpert pipeline.");
+      return signalPassFailure();
+    }
+    addMmt4dTilingExpertPassPipeline(pipeline, *maybeTilingConfig,
+                                     pipelineOpts);
     break;
   }
   case IREE::Codegen::DispatchLoweringPassPipeline::CPUDataTiling: {
-    TilingConfig tilingConfig = getTilingConfigForPipeline(funcOp);
-    addCPUDataTilingPipeline(pipeline, tilingConfig, pipelineOpts);
+    auto maybeTilingConfig = getTilingConfigForPipeline(funcOp);
+    if (failed(maybeTilingConfig)) {
+      funcOp.emitOpError(
+          "Tiling Config is necessary for CPUDataTiling pipeline.");
+      return signalPassFailure();
+    }
+    addCPUDataTilingPipeline(pipeline, *maybeTilingConfig, pipelineOpts);
     break;
   }
   case IREE::Codegen::DispatchLoweringPassPipeline::
       CPULinalgExtTileAndVectorize: {
-    TilingConfig tilingConfig = getTilingConfigForPipeline(funcOp);
-    addCPULinalgExtTileAndVectorizePipeline(pipeline, tilingConfig,
+    auto maybeTilingConfig = getTilingConfigForPipeline(funcOp);
+    if (failed(maybeTilingConfig)) {
+      funcOp.emitOpError("Tiling Config is necessary for "
+                         "CPULinalgExtTileAndVectorize pipeline.");
+      return signalPassFailure();
+    }
+    addCPULinalgExtTileAndVectorizePipeline(pipeline, *maybeTilingConfig,
                                             pipelineOpts);
     break;
   }

compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp

@@ -653,8 +653,14 @@ void addCPULinalgExtTileAndVectorizePipeline(
   }
 }
 
-void addCPUDefaultPassPipeline(OpPassManager &funcPassManager) {
+void addCPUDefaultPassPipeline(OpPassManager &funcPassManager,
+                               FailureOr<TilingConfig> &tilingConfig) {
   addTileAndDistributePasses(funcPassManager);
+  if (succeeded(tilingConfig) &&
+      tilingConfig.value().getNumTilingLevels() > 1) {
+    funcPassManager.addPass(createLLVMCPUTileAndFusePass(
+        tilingConfig.value().getVectorCommonParallelLevel()));
+  }
   addCPUBufferizePasses(funcPassManager);
 }
 

compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.h

@@ -96,10 +96,11 @@ void addCPULinalgExtTileAndVectorizePipeline(
     OpPassManager &funcPassManager, TilingConfig &tilingConfig,
     LLVMCPUPipelineOptions &pipelineOpt);
 
-/// Populates the passes to lower to scalars operations for linalg based
-/// code-generation. This pipeline does not vectorize, but instead just
-/// converts to memrefs
-void addCPUDefaultPassPipeline(OpPassManager &funcPassManager);
+/// Populates the passes to lower scalars and unknown tensor op (i.e. linalg op
+/// that is not specialized by any pipeline). Adds an additional level of tiling
+/// and converts to memrefs.
+void addCPUDefaultPassPipeline(OpPassManager &funcPassManager,
+                               FailureOr<TilingConfig> &tilingConfig);
 
 void addConvTileAndDecomposeExpertPassPipeline(
     OpPassManager &funcPassManager, TilingConfig &tilingConfig,

compiler/src/iree/compiler/DispatchCreation/FormDispatchRegions.cpp

@@ -547,14 +547,6 @@ isFusableWithConsumer(OpOperand &fusedOperand,
     return false;
   }
 
-  // TODO: Enable grouped convolution and depth wise pooling fusion.
-  // Rightnow, this is going through the default CPU pipeline and not through
-  // CONVTilingExpert.
-  if (isa<linalg::Conv2DNgchwFgchwOp, linalg::Conv2DNgchwGfchwOp,
-          linalg::PoolingNdhwcSumOp>(producer)) {
-    return false;
-  }
-
   auto producerFusionOp =
       dyn_cast<IREE::LinalgExt::LinalgFusionOpInterface>(producer);
   auto consumerFusionOp =