[CPU] Drop empty tile sizes from lowering config. (#21542)

If a tile size list has all zeros, it means no tiling at all. It is
redundant information in the IR. After switching to
`IREE::CPU::LoweringConfigAttr`, we no longer need to fill empty tile
sizes to make `TilingConfig` happy.

The revision deletes `emitInnerParallelList` option from
`LoweringConfigGenerator`, and it now only emits the tiling level if one
of the tile sizes is set.

The pipeline changes are needed because now the `rootOp` does not have
`vector_inner_parallel` level. It is mainly for the consumer ops that
have dimensions not captured by the `rootOp`. Given that it is the last
level of tiling, we move it out of the for loop.

---------

Signed-off-by: hanhanW <[email protected]>

Author: hanhanW

compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp

@@ -178,7 +178,8 @@ bool LoweringConfigAttr::hasTilingLevel(unsigned level) const {
 }
 
 bool LoweringConfigAttr::hasWorkgroupTilingLevel() const {
-  return !getWorkgroupTileSizes().empty();
+  return getConfig().contains(
+      getTilingLevelName(TilingLevel::DistributionTiles));
 }
 
 std::optional<unsigned> LoweringConfigAttr::getNumTilingLevels() const {

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

@@ -20,6 +20,7 @@
 #include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
 #include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
 #include "iree/compiler/Dialect/LinalgExt/Utils/IndexingUtils.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVectorExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/CommandLine.h"
@@ -110,6 +111,7 @@ static llvm::cl::opt<bool> clEnableRiscvAggressiveDist(
     llvm::cl::init(false));
 
 using IREE::Codegen::DispatchLoweringPassPipeline;
+using IREE::CPU::TilingLevel;
 
 // Encodes the pre-processing strategy to be applied on a Linalg operation
 // before vectorization.
@@ -960,14 +962,14 @@ static void setAlwaysVectorizeSizes(linalg::LinalgOp op,
 }
 
 /// A helper class to record different level tiling sizes and generate
-/// corresponding IREE::CPU::LoweringConfigAttr. Only vector level supports
-/// scalable tile sizes for now.
+/// corresponding IREE::CPU::LoweringConfigAttr for the rootOp. It implies that
+/// the distribution tiling level is always set, even if tile sizes are all
+/// zeros. Because a rootOp must have distribution tiling level.
+/// Only vector level supports scalable tile sizes for now.
 class LoweringConfigGenerator {
 public:
-  explicit LoweringConfigGenerator(Operation *op,
-                                   bool emitInnerParallelList = false)
-      : ctx(op->getContext()), rootOp(op),
-        emitInnerParallelList(emitInnerParallelList) {}
+  explicit LoweringConfigGenerator(Operation *op)
+      : ctx(op->getContext()), rootOp(op) {}
 
   void setDistributionTileSizes(ArrayRef<int64_t> tileSizes) {
     assert(distTileSizes.empty() && "expected to set only once");
@@ -990,7 +992,6 @@ class LoweringConfigGenerator {
   /// existing values. By default, it will always contain distribution tile
   /// sizes, unless the rootOp does not implement TilingInterface.
   IREE::CPU::LoweringConfigAttr generateCPULoweringConfig() {
-    using TilingLevel = IREE::CPU::TilingLevel;
     SmallVector<NamedAttribute> items;
     if (!distTileSizes.empty()) {
       appendLoweringConfigLevelAttr(items, TilingLevel::DistributionTiles,
@@ -1024,21 +1025,23 @@ class LoweringConfigGenerator {
                                     parallelTileSizes, parallelScalableFlags);
       appendLoweringConfigLevelAttr(items, TilingLevel::VectorReductionTiles,
                                     reductionTileSizes, reductionScalableFlags);
-      if (emitInnerParallelList) {
-        size_t size = parallelTileSizes.size();
-        appendLoweringConfigLevelAttr(
-            items, TilingLevel::VectorInnerParallelTiles,
-            SmallVector<int64_t>(size, 0), SmallVector<bool>(size, false));
-      }
     }
     return IREE::CPU::LoweringConfigAttr::get(ctx, items);
   }
 
 private:
+  /// Appends the `level` with (`tileSizes`, `scalableFlags`) tiling config to
+  /// `items`, if it is not a NOP config. E.g., if all the tile sizes are zeros,
+  /// it means no tiling at all. Only the distribution tiling level is
+  /// unconditionally added because a root op expects the level to be present.
   void appendLoweringConfigLevelAttr(SmallVectorImpl<NamedAttribute> &items,
-                                     IREE::CPU::TilingLevel level,
+                                     TilingLevel level,
                                      ArrayRef<int64_t> tileSizes,
                                      ArrayRef<bool> scalableFlags = {}) {
+    if (level != TilingLevel::DistributionTiles &&
+        llvm::all_of(tileSizes, [](int64_t v) { return v == 0; })) {
+      return;
+    }
     items.emplace_back(IREE::CPU::getTilingLevelName(level),
                        IREE::CPU::LoweringConfigAttr::getTilingLevelAttr(
                            ctx, tileSizes, scalableFlags));
@@ -1047,12 +1050,6 @@ class LoweringConfigGenerator {
   MLIRContext *ctx;
   Operation *rootOp;
 
-  // Generates the `IREE::CPU::TilingLevel::VectorInnerParallelTiles` tile sizes
-  // in the lowering config. Usually, they are zero values.
-  // TODO(hanchung): Remove the field once all the pipelines switch to CPU
-  // lowering_config. It is alive for legacy setup.
-  bool emitInnerParallelList = false;
-
   // The tile sizes for distribution from the `rootOp`'s perspective.
   SmallVector<int64_t> distTileSizes;
 
@@ -1092,8 +1089,14 @@ static IREE::Codegen::LoweringConfigAttrInterface getNewLoweringConfig(
 
   SmallVector<NamedAttribute> newItems;
   for (auto [level, tileSizes, scalableFlags] : tilingInfo) {
-    if (!setDistributionConfig &&
-        level == IREE::CPU::TilingLevel::DistributionTiles) {
+    if (!setDistributionConfig && level == TilingLevel::DistributionTiles) {
+      continue;
+    }
+    // Distribution tile sizes is a must for rootOp, because it is the
+    // definition of root op. An operation that has distribution tile sizes is
+    // the root op. Other level can be dropped if all the tile sizes are zeros.
+    if (level != TilingLevel::DistributionTiles &&
+        llvm::all_of(tileSizes, [](int64_t val) { return val == 0; })) {
       continue;
     }
     newItems.emplace_back(IREE::CPU::getTilingLevelName(level),
@@ -1155,7 +1158,7 @@ static LogicalResult setMatmulPeelingRootConfig(
                                         inputVecScalableTileFlags.end());
   vectorScalableFlags.back() = false;
 
-  LoweringConfigGenerator generator(op, /*emitInnerParallelList=*/true);
+  LoweringConfigGenerator generator(op);
   generator.setDistributionTileSizes(distTileSizes);
   generator.setCacheTileSizes(cacheTileSizes);
   generator.setVectorTileSizes(vecTileSizes, vectorScalableFlags);
@@ -1206,7 +1209,7 @@ static LogicalResult setMatmulRootConfig(
   }
   limitVectorTileSizes(cast<linalg::LinalgOp>(op.getOperation()), vecTileSizes);
 
-  LoweringConfigGenerator generator(op, /*emitInnerParallelList=*/true);
+  LoweringConfigGenerator generator(op);
   generator.setDistributionTileSizes(distTileSizes);
   generator.setVectorTileSizes(vecTileSizes, vecScalableFlags);
   IREE::CPU::LoweringConfigAttr loweringConfig =
@@ -2085,10 +2088,9 @@ setDefaultGenericOpRootConfig(mlir::FunctionOpInterface entryPointFn,
   // If there are no loops, there is nothing to do.
   unsigned numLoops = genericOp.getNumLoops();
   if (numLoops == 0) {
+    LoweringConfigGenerator generator(genericOp);
     return setOpConfigAndEntryPointFnTranslation(
-        entryPointFn, genericOp,
-        IREE::CPU::LoweringConfigAttr::get(genericOp.getContext(),
-                                           SmallVector<NamedAttribute>()),
+        entryPointFn, genericOp, generator.generateCPULoweringConfig(),
         DispatchLoweringPassPipeline::CPUDefault);
   }
 
@@ -2113,7 +2115,7 @@ setDefaultGenericOpRootConfig(mlir::FunctionOpInterface entryPointFn,
                      distConfig.maxTileSizes, vecPreProcStrategy, vecTileSizes);
   limitVectorTileSizes(genericOp, vecTileSizes);
 
-  LoweringConfigGenerator generator(genericOp, /*emitInnerParallelList=*/true);
+  LoweringConfigGenerator generator(genericOp);
   generator.setDistributionTileSizes(distTileSizes);
   generator.setVectorTileSizes(vecTileSizes);
   IREE::CPU::LoweringConfigAttr loweringConfig =
@@ -2267,7 +2269,7 @@ setTransposeLikeOpRootConfig(mlir::FunctionOpInterface entryPointFn,
   SmallVector<int64_t> distTileSizes =
       getDefaultDistributedLevelTileSizes(genericOp, distConfig);
 
-  LoweringConfigGenerator generator(genericOp, /*emitInnerParallelList=*/true);
+  LoweringConfigGenerator generator(genericOp);
   generator.setDistributionTileSizes(distTileSizes);
   generator.setVectorTileSizes(vecSizes, vecScalableDims);
   IREE::CPU::LoweringConfigAttr loweringConfig =
@@ -2346,7 +2348,7 @@ static LogicalResult setElementwiseGenericOpRootConfig(
                       vecPreProcStrategy == VectorPreProcStrategy::Masking);
   }
 
-  LoweringConfigGenerator generator(genericOp, /*emitInnerParallelList=*/true);
+  LoweringConfigGenerator generator(genericOp);
   generator.setDistributionTileSizes(distTileSizes);
   generator.setVectorTileSizes(vecTileSizes);
   IREE::CPU::LoweringConfigAttr loweringConfig =
@@ -2603,7 +2605,7 @@ static LogicalResult setRootConfig(mlir::FunctionOpInterface entryPointFn,
 
   SmallVector<int64_t> distTileSizes =
       getDefaultDistributedLevelTileSizes(padOp, distConfig);
-  LoweringConfigGenerator generator(padOp, /*emitInnerParallelList=*/true);
+  LoweringConfigGenerator generator(padOp);
   generator.setDistributionTileSizes(distTileSizes);
   generator.setVectorTileSizes(distConfig.vectorSizeHints);
   IREE::CPU::LoweringConfigAttr loweringConfig =
@@ -2879,7 +2881,7 @@ adjustTileSizesForGenericOp(mlir::FunctionOpInterface entryPointFn,
 /// `level`, if it is present. Otherwise, adds a new item to the vector.
 static void updateOrAddTilingLevelInfo(
     SmallVectorImpl<IREE::CPU::LoweringConfigLevelInfo> &tilingInfo,
-    IREE::CPU::TilingLevel level, ArrayRef<int64_t> tileSizes,
+    TilingLevel level, ArrayRef<int64_t> tileSizes,
     ArrayRef<bool> scalableFlags) {
   for (IREE::CPU::LoweringConfigLevelInfo &info : tilingInfo) {
     if (info.level == level) {
@@ -2946,16 +2948,18 @@ setLoweringConfigForComputeOps(mlir::FunctionOpInterface entryPointFn,
   }
 
   auto rootLoweringConfig = getLoweringConfig(rootOperation);
-  std::unique_ptr<TilingConfig> tilingConfig =
-      TilingConfig::create(rootLoweringConfig);
   SmallVector<int64_t> distTileSizes, parallelVecTileSizes;
   SmallVector<bool> distScalableTileSizes, parallelVecScalableTileSizes;
-  if (tilingConfig->getNumTilingLevels() > 0) {
-    distTileSizes = tilingConfig->getDistributionTileSizes();
-  }
-  if (tilingConfig->getNumTilingLevels() > 1) {
-    std::tie(parallelVecTileSizes, parallelVecScalableTileSizes) =
-        tilingConfig->getVectorCommonParallelSizes();
+  assert(rootLoweringConfig.hasWorkgroupTilingLevel());
+  distTileSizes = rootLoweringConfig.getWorkgroupTileSizes();
+  if (rootLoweringConfig.hasTilingLevel(
+          TilingLevel::VectorCommonParallelTiles)) {
+    auto attr = cast<IREE::Codegen::LoweringConfigTilingLevelAttr>(
+        rootLoweringConfig.getTilingLevelAttr(
+            TilingLevel::VectorCommonParallelTiles));
+    parallelVecTileSizes.assign(attr.getSizes().begin(), attr.getSizes().end());
+    parallelVecScalableTileSizes.assign(attr.getScalableFlags().begin(),
+                                        attr.getScalableFlags().end());
   }
 
   size_t maxLoopNums = 0;
@@ -3080,6 +3084,9 @@ setLoweringConfigForComputeOps(mlir::FunctionOpInterface entryPointFn,
   }
 
   // Set the lowering configs with new tile sizes.
+  // TODO(hanchung): Deprecate TilingConfig from the file.
+  std::unique_ptr<TilingConfig> tilingConfig =
+      TilingConfig::create(rootLoweringConfig);
   for (auto op : computeOps) {
     int numLoops = cast<TilingInterface>(op).getLoopIteratorTypes().size();
     SmallVector<IREE::CPU::LoweringConfigLevelInfo> newTilingInfo;

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

@@ -269,6 +269,10 @@ void LLVMCPUTileAndFusePass::runOnOperation() {
     LDBG() << "can't find lowering_config, skip TileAndFuse";
     return;
   }
+  if (!loweringConfig.hasTilingLevel(tilingLevel)) {
+    LDBG() << "no tile sizes for the tiling level, skip TileAndFuse";
+    return;
+  }
 
   SmallVector<int64_t> tileSizes;
   SmallVector<bool> tileScalableFlags;

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

@@ -434,18 +434,24 @@ void addMultiTilingExpertPassPipeline(OpPassManager &funcPassManager,
           static_cast<IREE::CPU::TilingLevel>(i), /*skipRootOp=*/true));
       break;
     case IREE::CPU::TilingLevel::VectorInnerParallelTiles:
-      funcPassManager.addPass(createLLVMCPUTileAndFusePass(
-          tilingConfig.getVectorInnerParallelLevel()));
-      break;
     case IREE::CPU::TilingLevel::DistributionTiles:
     case IREE::CPU::TilingLevel::MaxNumTileLevels:
     case IREE::CPU::TilingLevel::InvalidLevel:
-      break;
+      continue;
     };
     funcPassManager.addPass(createFuseTensorPadWithConsumerPass());
     funcPassManager.addPass(createConcretizePadResultShapePass());
   }
 
+  // `VectorInnerParallelTiles` level models the tiling and fusion for the
+  // dimensions that are not captured in root op. I.e., root op may not have the
+  // config for the level. Thus, we run the LLVMCPUTileAndFuse pass for
+  // consumers.
+  funcPassManager.addPass(createLLVMCPUTileAndFusePass(
+      IREE::CPU::TilingLevel::VectorInnerParallelTiles));
+  funcPassManager.addPass(createFuseTensorPadWithConsumerPass());
+  funcPassManager.addPass(createConcretizePadResultShapePass());
+
   funcPassManager.addPass(createForallToForPass());
   if (pipelineOpt.enablePeeling) {
     funcPassManager.addPass(createLLVMCPUPeelPass());

compiler/src/iree/compiler/Codegen/LLVMCPU/test/2d_scalable_to_1d_scalable.mlir

@@ -1,7 +1,7 @@
 // RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-llvmcpu-2d-scalable-to-1d-scalable{assume-arm-sme=true},cse))" --split-input-file %s | FileCheck %s
 
-#compute_config = #iree_cpu.lowering_config<vector_common_parallel = [[4], [4]], vector_reduction = [0, 0], vector_inner_parallel = [0, 0]>
-#matmul_config = #iree_cpu.lowering_config<distribution = [0, 0, 0], vector_common_parallel =  [[4], [4], 0], vector_reduction = [0, 0, 1], vector_inner_parallel = [0, 0, 0]>
+#compute_config = #iree_cpu.lowering_config<vector_common_parallel = [[4], [4]]>
+#matmul_config = #iree_cpu.lowering_config<distribution = [0, 0, 0], vector_common_parallel =  [[4], [4], 0], vector_reduction = [0, 0, 1]>
 #dim_0_map = affine_map<(d0)[s0] -> (-d0 + 32400, s0)>
 #dim_1_map = affine_map<(d0)[s0] -> (-d0 + 16, s0)>
 
@@ -60,9 +60,9 @@ func.func @scalable_2d_matmul_and_generic(%arg0: tensor<32400x32xf32>, %arg1: te
   }
   return %2 : tensor<32400x16xf32>
 }
-// CHECK: #[[FILL_CONFIG:.*]] = #iree_cpu.lowering_config<vector_common_parallel = {{\[}}[4], [4]], vector_inner_parallel = [0, 0], vector_reduction = [0, 0]>
-// CHECK: #[[MATMUL_CONFIG:.*]] = #iree_cpu.lowering_config<distribution = [0, 0, 0], vector_common_parallel = {{\[}}[4], [4], 0], vector_inner_parallel = [0, 0, 0], vector_reduction = [0, 0, 1]>
-// CHECK: #[[GENERIC_CONFIG:.*]] = #iree_cpu.lowering_config<vector_common_parallel = [4, [4]], vector_inner_parallel = [0, 0], vector_reduction = [0, 0]>
+// CHECK: #[[FILL_CONFIG:.*]] = #iree_cpu.lowering_config<vector_common_parallel = {{\[}}[4], [4]]>
+// CHECK: #[[MATMUL_CONFIG:.*]] = #iree_cpu.lowering_config<distribution = [0, 0, 0], vector_common_parallel = {{\[}}[4], [4], 0], vector_reduction = [0, 0, 1]>
+// CHECK: #[[GENERIC_CONFIG:.*]] = #iree_cpu.lowering_config<vector_common_parallel = [4, [4]]>
 //
 //      CHECK: func.func @scalable_2d_matmul_and_generi
 //      CHECK:   %[[C4:.*]] = arith.constant 4 : index

compiler/src/iree/compiler/Codegen/LLVMCPU/test/pipeline_arm_sme_streaming_mode_tests.mlir

@@ -6,7 +6,7 @@
   #hal.pipeline.binding<storage_buffer>,
   #hal.pipeline.binding<storage_buffer>
 ]>
-#config = #iree_cpu.lowering_config<distribution = [0], vector_common_parallel = [1], vector_reduction = [0], vector_inner_parallel = [0]>
+#config = #iree_cpu.lowering_config<distribution = [0], vector_common_parallel = [1]>
 module {
 module {
   func.func @fixed_size_dispatch() attributes {hal.executable.target = #hal.executable.target<"llvm-cpu", "embedded-elf-arm_64", {cpu_features = "+sve,+sme", data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", native_vector_size = 16 : index, target_triple = "aarch64-none-elf"}>,
@@ -43,7 +43,7 @@ module {
   #hal.pipeline.binding<storage_buffer>,
   #hal.pipeline.binding<storage_buffer>
 ]>
-#config = #iree_cpu.lowering_config<distribution = [0], vector_common_parallel = [[1]], vector_reduction = [0], vector_inner_parallel = [0]>
+#config = #iree_cpu.lowering_config<distribution = [0], vector_common_parallel = [[1]]>
 module {
 module {
   func.func @scalable_dispatch() attributes {hal.executable.target = #hal.executable.target<"llvm-cpu", "embedded-elf-arm_64", {cpu_features = "+sve,+sme", data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", native_vector_size = 16 : index, target_triple = "aarch64-none-elf"}>,
@@ -81,7 +81,7 @@ module {
   #hal.pipeline.binding<storage_buffer>,
   #hal.pipeline.binding<storage_buffer>
 ]>
-#config = #iree_cpu.lowering_config<distribution = [0, 0], vector_common_parallel = [[4], [4]], vector_reduction = [0, 0], vector_inner_parallel = [0, 0]>
+#config = #iree_cpu.lowering_config<distribution = [0, 0], vector_common_parallel = [[4], [4]]>
 module {
 module {
   func.func @scalable_dispatch_using_za() attributes {hal.executable.target = #hal.executable.target<"llvm-cpu", "embedded-elf-arm_64", {cpu_features = "+sve,+sme", data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", native_vector_size = 16 : index, target_triple = "aarch64-none-elf"}>,