[CPU] Switch CPUDoubleTilingExpert pipeline to use IREE::CPU::LoweringConfigAttr. (iree-org#21354)

The revision switches all the dispatches that use
`CPUDoubleTilingExpert` to `IREE::CPU::LoweringConfigAttr`, which is a
root-based tiling approach. There are two commits in the revision:

- [The switch for `linalg.matmul`
dispatches](iree-org@e885e86):
it mainly focuses on the lowering config changes.
- [The switch for `linalg.generic`
dispatches](iree-org@48a527f):
  - Changes for lowering configs.
  - Update the pipeline that enumerates all the tiling levels.
- Update `LLVMCPU2DScalableTo1DScalable` pass to use
`IREE::CPU::LoweringConfigAttr`.

The pipeline and lowering config verification only applies on the root
op because CPUDoubleTilingExpert expects exactly four levels of tiling.
After the switch, only the root op has four levels of tiling. We will
need to refresh the verification logic anyway, as they are legacy code
and it is easier to have better implementation today. So it will be
refreshed in a follow-up.

New known issues:
- Trailing vector unit dims are not folded away in the root-op based
pipeline, which results in larger binary sizes. Because they are
unrolled: iree-org#21420

---------

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

Author: hanhanW

compiler/src/iree/compiler/Codegen/Common/TileSizeSelection.cpp

@@ -8,8 +8,10 @@
 #include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenInterfaces.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
+#include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinAttributes.h"
 
 #define DEBUG_TYPE "tiling-config"
@@ -169,88 +171,57 @@ SizesAndScalableFlags TilingConfig::getVectorTileSizes() {
   return std::make_pair(vectorSizes, scalableFlags);
 }
 
-/// Returns a new `LoweringConfigAttr`, with the tile sizes of vector
-/// dimensions, set to `sizes`, and the corresponding scalability set to
-/// `scalableFlags`.
-IREE::Codegen::LoweringConfigAttr
-TilingConfig::getLoweringConfigWithNewVectorSizes(
+IREE::CPU::LoweringConfigAttr TilingConfig::getLoweringConfigWithNewVectorSizes(
     ArrayRef<int64_t> sizes, ArrayRef<bool> scalableFlags) {
   unsigned numDims = getNumDimensions();
+  (void)numDims;
   assert(sizes.size() == numDims &&
          "expected `sizes` to match number of dimensions");
   assert((scalableFlags.empty() || scalableFlags.size() == numDims) &&
          "expected `scalableFlags` to match "
          "number of dimensions (or be empty)");
 
-  // Make a map from tiling levels to vector dims at that level.
-  std::array<SmallVector<unsigned, 4>, TilingLevel::MaxNumTileLevels>
-      tilingLevelToDimsMap;
-  for (unsigned dimPos = 0; dimPos < numDims; ++dimPos) {
-    auto tilingLevelIndex = getTilingLevelForVectorDimPosition(dimPos);
-    assert((tilingLevelIndex.has_value() || sizes[dimPos] == 0) &&
-           "attempting to set vector size for dim with underspecified tiling "
-           "level (zero is the only valid tile size)");
-    if (tilingLevelIndex.has_value())
-      tilingLevelToDimsMap[*tilingLevelIndex].push_back(dimPos);
-  }
-
-  MLIRContext *context = loweringConfig.getContext();
-  SmallVector<IREE::Codegen::LoweringConfigTilingLevelAttr> tilingLevels;
-  for (unsigned i = 0, e = getNumTilingLevels(); i < e; ++i) {
-    tilingLevels.push_back(cast<IREE::Codegen::LoweringConfigTilingLevelAttr>(
-        loweringConfig.getTilingLevelAttr(i)));
-  }
-  SmallVector<IREE::Codegen::LoweringConfigTilingLevelAttr> newTilingLevelsList(
-      tilingLevels.begin(), tilingLevels.end());
-
-  // For each vector tiling level:
-  for (auto [tilingLevelIndex, tilingLevelDims] :
-       llvm::enumerate(tilingLevelToDimsMap)) {
-    if (tilingLevelDims.empty())
+  MLIRContext *ctx = loweringConfig.getContext();
+  SmallVector<NamedAttribute> items;
+  for (unsigned i = 0, e = TilingLevel::MaxNumTileLevels; i < e; ++i) {
+    auto level = static_cast<TilingLevel>(i);
+    if (!isValidLevel(level)) {
       continue;
-    auto level = tilingLevels[tilingLevelIndex];
-    SmallVector<int64_t> newSizes(level.getSizes());
-    SmallVector<bool> newScalableFlags(level.getScalableFlags());
-    newScalableFlags.resize(numDims);
-    // 1. Update all the vector sizes within that tiling level.
-    for (unsigned dimPos : tilingLevelDims) {
-      std::tie(newSizes[dimPos], newScalableFlags[dimPos]) =
-          getTileSizeAtIndex(sizes, scalableFlags, dimPos);
     }
-    // 2. Then create a new tiling level attribute for that level.
-    auto newLevel = IREE::Codegen::LoweringConfigTilingLevelAttr::get(
-        context, newSizes, level.getInterchange(), newScalableFlags);
-    newTilingLevelsList[tilingLevelIndex] = newLevel;
-  }
-
-  // Create a new `lowering_config` attribute.
-  auto newTilingLevels = IREE::Codegen::LoweringConfigTilingLevelsAttr::get(
-      context, newTilingLevelsList);
-  return IREE::Codegen::LoweringConfigAttr::get(context, newTilingLevels);
-}
-
-/// Returns a list with the tiling levels that can be fused for this
-/// configuration.
-SmallVector<int64_t> TilingConfig::getFusableLevels() {
-  switch (getNumTilingLevels()) {
-  case 0:
-    return {};
-  case 1:
-    // Only distribution level.
-    return {0};
-  case 3:
-    // Only distribution level + vector common parallel levels.
-    return {0, 1};
-  case 4:
-    // Distribution + vector common parallel levels + vector inner parallel
-    // levels.
-    return {0, 1, 3};
-  case 6:
-    // Distribution + cache parallel levels.
-    return {0, 1, 3, 5};
-  default:
-    llvm_unreachable("Unexpected number of tiling levels");
+    switch (level) {
+    case TilingLevel::DistributionTiles:
+    case TilingLevel::CacheParallelTiles:
+    case TilingLevel::CacheReductionTiles: {
+      items.emplace_back(IREE::CPU::getTilingLevelName(level),
+                         getTilingLevelAttr(i));
+      break;
+    }
+    case TilingLevel::VectorCommonParallelTiles:
+    case TilingLevel::VectorReductionTiles:
+    case TilingLevel::VectorInnerParallelTiles: {
+      auto attr = cast<IREE::Codegen::LoweringConfigTilingLevelAttr>(
+          loweringConfig.getTilingLevelAttr(i));
+      SmallVector<int64_t> newSizes(attr.getSizes());
+      SmallVector<bool> newScalableFlags(attr.getScalableFlags());
+      newScalableFlags.resize(newSizes.size(), false);
+      for (auto [idx, size] : llvm::enumerate(newSizes)) {
+        if (size == 0) {
+          continue;
+        }
+        newSizes[idx] = sizes[idx];
+        newScalableFlags[idx] = scalableFlags[idx];
+      }
+      auto newLevel = IREE::Codegen::LoweringConfigTilingLevelAttr::get(
+          ctx, newSizes, attr.getInterchange(), newScalableFlags);
+      items.emplace_back(IREE::CPU::getTilingLevelName(level), newLevel);
+      break;
+    }
+    case TilingLevel::MaxNumTileLevels:
+    case TilingLevel::InvalidLevel:
+      break;
+    };
   }
+  return IREE::CPU::LoweringConfigAttr::get(ctx, items);
 }
 
 /// Returns the actual level in the configuration for this level of tiling.

compiler/src/iree/compiler/Codegen/Common/TileSizeSelection.h

@@ -181,19 +181,26 @@ class TilingConfig {
   /// Returns a new `LoweringConfigAttr`, with the tile sizes of vector
   /// dimensions, set to `sizes`, and the corresponding scalability set to
   /// `scalableFlags`.
-  IREE::Codegen::LoweringConfigAttr
+  IREE::CPU::LoweringConfigAttr
   getLoweringConfigWithNewVectorSizes(ArrayRef<int64_t> sizes,
                                       ArrayRef<bool> scalableFlags = {});
 
-  /// Returns a list with the tiling levels that can be fused for this
-  /// configuration.
-  SmallVector<int64_t> getFusableLevels();
-
-  // TODO(dcaballe): Revisit if these features are ever used.
-  SmallVector<int64_t> getTileInterchangeSizes(unsigned level) {
-    auto attr = cast<IREE::Codegen::LoweringConfigTilingLevelAttr>(
-        loweringConfig.getTilingLevelAttr(level));
-    return SmallVector<int64_t>(attr.getInterchange());
+  /// Returns the `level`-th valid tiling attribute. Returns an empty vector if
+  /// it does not exist.
+  IREE::Codegen::LoweringConfigTilingLevelAttr
+  getTilingLevelAttr(int64_t level) {
+    for (auto [idx, mappedLevel] :
+         llvm::enumerate(tilingLevelToActualLevelMap)) {
+      if (mappedLevel == TilingLevel::InvalidLevel) {
+        continue;
+      }
+      if (--level >= 0) {
+        continue;
+      }
+      return cast<IREE::Codegen::LoweringConfigTilingLevelAttr>(
+          loweringConfig.getTilingLevelAttr(mappedLevel));
+    }
+    return {};
   }
 
 private:

compiler/src/iree/compiler/Codegen/LLVMCPU/BUILD.bazel

@@ -166,6 +166,7 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:SCFUtils",
         "@llvm-project//mlir:TensorDialect",
         "@llvm-project//mlir:TensorTransforms",
+        "@llvm-project//mlir:TilingInterface",
         "@llvm-project//mlir:TosaDialect",
         "@llvm-project//mlir:TosaToArith",
         "@llvm-project//mlir:TransformDialect",

compiler/src/iree/compiler/Codegen/LLVMCPU/CMakeLists.txt

@@ -135,6 +135,7 @@ iree_cc_library(
     MLIRSCFUtils
     MLIRTensorDialect
     MLIRTensorTransforms
+    MLIRTilingInterface
     MLIRTosaDialect
     MLIRTosaToArith
     MLIRTransformDialect