[LLVMCPU] Add additional level of tiling to the default

iree-org · Nov 10, 2024 · 864d9df · 864d9df
1 parent 5c45591
commit 864d9df
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 8 deletions.
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
@@ -2508,6 +2508,12 @@ static LogicalResult setRootConfig(mlir::FunctionOpInterface entryPointFn,
   SmallVector<int64_t> distTileSizes =
       getDefaultDistributedLevelTileSizes(op, DistributionHeuristicConfig{});
   TileSizesListType tileSizes = {distTileSizes};
+  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);
 }

diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPULowerExecutableTarget.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPULowerExecutableTarget.cpp
@@ -84,6 +84,7 @@ getRootLoweringConfig(FunctionOpInterface funcOp) {
 
 static TilingConfig getTilingConfigForPipeline(FunctionOpInterface funcOp) {
   auto maybeLoweringConfig = getRootLoweringConfig(funcOp);
+  llvm::errs()<<"Hey I am here";
   assert(succeeded(maybeLoweringConfig) &&
          "Pipeline requires a lowering config");
   return TilingConfig(*maybeLoweringConfig);
@@ -122,9 +123,11 @@ 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: {
+    TilingConfig tilingConfig = getTilingConfigForPipeline(funcOp);
+    addCPUDefaultPassPipeline(pipeline, tilingConfig);
     break;
+    }
   case IREE::Codegen::DispatchLoweringPassPipeline::
       CPUBufferOpsTileAndVectorize: {
     TilingConfig tilingConfig = getTilingConfigForPipeline(funcOp);

diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp
@@ -653,8 +653,12 @@ void addCPULinalgExtTileAndVectorizePipeline(
   }
 }
 
-void addCPUDefaultPassPipeline(OpPassManager &funcPassManager) {
+void addCPUDefaultPassPipeline(OpPassManager &funcPassManager, TilingConfig &tilingConfig) {
   addTileAndDistributePasses(funcPassManager);
+  if(tilingConfig.getNumTilingLevels() > 1){
+  funcPassManager.addPass(createLLVMCPUTileAndFusePass(
+      tilingConfig.getVectorCommonParallelLevel()));
+  }
   addCPUBufferizePasses(funcPassManager);
 }
 

diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.h b/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.h
@@ -99,7 +99,7 @@ void addCPULinalgExtTileAndVectorizePipeline(
 /// 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);
+void addCPUDefaultPassPipeline(OpPassManager &funcPassManager, TilingConfig &tilingConfig);
 
 void addConvTileAndDecomposeExpertPassPipeline(
     OpPassManager &funcPassManager, TilingConfig &tilingConfig,

diff --git a/compiler/src/iree/compiler/DispatchCreation/FormDispatchRegions.cpp b/compiler/src/iree/compiler/DispatchCreation/FormDispatchRegions.cpp
@@ -550,10 +550,10 @@ isFusableWithConsumer(OpOperand &fusedOperand,
   // 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;
-  }
+  // if (isa<linalg::Conv2DNgchwFgchwOp, linalg::Conv2DNgchwGfchwOp,
+  //         linalg::PoolingNdhwcSumOp>(producer)) {
+  //   return false;
+  // }
 
   auto producerFusionOp =
       dyn_cast<IREE::LinalgExt::LinalgFusionOpInterface>(producer);