[CPU] Drop TilingConfig from KernelDispatch.cpp (#21567)

The revision implements `getAvailableTilingInfo` in
`IREE::CPU::LoweringConfigAttr`, and replace all the uses of
`TilingConfig`'s version with it. It also deletes the method from
`TilingConfig` and makes corresponding changes to the codebase.

The lowering config propagation is disabled if the root op does not use
`IREE::CPU::LoweringConfigAttr`, because we need the method for the
propagation. However, there is still a case that it uses Codegen
lowering config, which is `LinalgExt::CustomOp` dispatch. The zeros in
tiling config are dropped in the propagation. Thus, there is a change in
lit test. It is a reasonable change, because it keeps the same behavior
as GPU backends. See
ab88871

---------

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

Author: hanhanW

compiler/src/iree/compiler/Codegen/Common/CPU/CPUPrepareUkernels.cpp

@@ -151,11 +151,10 @@ static LogicalResult reduceDefiningOp(PatternRewriter &rewriter, Value input) {
 
 /// Drops the first element from all the tile sizes list. The first element is
 /// for the batch dimension.
-static IREE::Codegen::LoweringConfigAttrInterface
-dropBatchTileSize(IREE::Codegen::LoweringConfigAttrInterface config) {
-  std::unique_ptr<TilingConfig> tilingConfig = TilingConfig::create(config);
+static IREE::CPU::LoweringConfigAttr
+dropBatchTileSize(IREE::CPU::LoweringConfigAttr config) {
   SmallVector<IREE::CPU::LoweringConfigLevelInfo> tilingInfo =
-      tilingConfig->getTilingLevelInfo();
+      config.getAvailableTilingInfo();
   SmallVector<NamedAttribute> newItems;
   for (auto [level, tileSizes, scalableTileFlags] : tilingInfo) {
     tileSizes.erase(tileSizes.begin());
@@ -204,8 +203,7 @@ struct ConvertBatchMmt4DtoMmt4DPattern
               .result();
 
       auto loweringConfig =
-          getLoweringConfig<IREE::Codegen::LoweringConfigAttrInterface>(
-              oldFillOp);
+          getLoweringConfig<IREE::CPU::LoweringConfigAttr>(oldFillOp);
       if (loweringConfig) {
         auto config = dropBatchTileSize(loweringConfig);
         setLoweringConfig(reducedOut.getDefiningOp(), config);
@@ -240,8 +238,7 @@ struct ConvertBatchMmt4DtoMmt4DPattern
         loc, reducedOut.getType(), ValueRange{reducedLhs, reducedRhs},
         ValueRange{reducedOut});
 
-    auto loweringConfig =
-        getLoweringConfig<IREE::Codegen::LoweringConfigAttrInterface>(op);
+    auto loweringConfig = getLoweringConfig<IREE::CPU::LoweringConfigAttr>(op);
     if (loweringConfig) {
       auto config = dropBatchTileSize(loweringConfig);
       setLoweringConfig(mmt4DOp, config);

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

@@ -95,24 +95,6 @@ void TilingConfig::initFromCPULoweringConfig(IREE::CPU::LoweringConfigAttr lc) {
   }
 }
 
-SmallVector<IREE::CPU::LoweringConfigLevelInfo>
-TilingConfig::getTilingLevelInfo() {
-  SmallVector<IREE::CPU::LoweringConfigLevelInfo> result;
-  TileSizesListType tileSizesList = getTileSizes();
-  ScalableTileFlagsListType scalableFlagsList = getScalableTileFlags();
-  int64_t mappedIdx = 0;
-  for (auto [idx, actualLevel] : llvm::enumerate(tilingLevelToActualLevelMap)) {
-    if (actualLevel == IREE::CPU::TilingLevel::InvalidLevel) {
-      continue;
-    }
-    result.push_back(IREE::CPU::LoweringConfigLevelInfo{
-        static_cast<IREE::CPU::TilingLevel>(idx), tileSizesList[mappedIdx],
-        scalableFlagsList[mappedIdx]});
-    mappedIdx++;
-  }
-  return result;
-}
-
 bool TilingConfig::isValidLevel(IREE::CPU::TilingLevel level) {
   return tilingLevelToActualLevelMap[static_cast<int64_t>(level)] !=
          IREE::CPU::TilingLevel::InvalidLevel;

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

@@ -69,12 +69,6 @@ class TilingConfig {
                           [](int64_t tileSize) { return tileSize != 0; });
   }
 
-  /// Returns a list of tiling information for each level. Each value is a valid
-  /// level in the TilingConfig.
-  /// Different from attribute variant, the method materialize the attribute
-  /// content to the `IREE::CPU::LoweringConfigLevelInfo` contrainer.
-  SmallVector<IREE::CPU::LoweringConfigLevelInfo> getTilingLevelInfo();
-
   /// Returns all the tile sizes of all the levels of the configuration.
   TileSizesListType getTileSizes() const {
     TileSizesListType result;

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

@@ -158,6 +158,24 @@ Attribute LoweringConfigAttr::getTilingLevelAttr(MLIRContext *ctx,
       ctx, tileSizes, /*interchange=*/{}, scalableFlags);
 }
 
+SmallVector<LoweringConfigLevelInfo>
+LoweringConfigAttr::getAvailableTilingInfo() {
+  SmallVector<LoweringConfigLevelInfo> result;
+  for (unsigned i = 0, e = TilingLevel::MaxNumTileLevels; i < e; ++i) {
+    if (!hasTilingLevel(i)) {
+      continue;
+    }
+    auto attr = cast<IREE::Codegen::LoweringConfigTilingLevelAttr>(
+        getTilingLevelAttr(i));
+    LoweringConfigLevelInfo item;
+    item.level = static_cast<TilingLevel>(i);
+    llvm::append_range(item.sizes, attr.getSizes());
+    llvm::append_range(item.scalableFlags, attr.getScalableFlags());
+    result.push_back(item);
+  }
+  return result;
+}
+
 SmallVector<int64_t> LoweringConfigAttr::getWorkgroupTileSizes() const {
   return getTileSizes(getConfig(), DistributionTiles);
 }

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

@@ -71,6 +71,9 @@ def IREECPU_LoweringConfigAttr :
     static Attribute getTilingLevelAttr(MLIRContext *ctx,
                                    ArrayRef<int64_t> tileSizes,
                                    ArrayRef<bool> scalableFlags);
+
+    /// Returns a vector that contains all the tiling information in the config.
+    SmallVector<LoweringConfigLevelInfo> getAvailableTilingInfo();
   }];
   let hasCustomAssemblyFormat = 1;
   let genVerifyDecl = 1;

compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h

@@ -11,12 +11,6 @@
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenInterfaces.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h"
 
-// clang-format off
-#define GET_ATTRDEF_CLASSES
-#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.h.inc"
-#undef GET_ATTRDEF_CLASSES
-// clang-format on
-
 namespace mlir::iree_compiler::IREE::CPU {
 
 /// Representation for all the supported tiling levels. All or just a subset of
@@ -43,4 +37,8 @@ StringRef getTilingLevelName(TilingLevel level);
 
 } // namespace mlir::iree_compiler::IREE::CPU
 
+// clang-format off
+#define GET_ATTRDEF_CLASSES
+#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.h.inc"
+// clang-format on
 #endif // IREE_COMPILER_CODEGEN_DIALECT_CPU_IREECPUTYPES_H_

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

@@ -1063,30 +1063,12 @@ class LoweringConfigGenerator {
 };
 
 /// Returns the same lowering_config attribute with the updated tile sizes and
-/// scalable tile flags. The `setDistrubtionConfig` flag is only available when
-/// `origLoweringConfig is a IREE::CPU::LoweringConfigAttr. The distribution
-/// tiling sizes is not set if it is false.
-/// See `Codegen/Common/TileSizeSelection.h` for the convention of mapping
-/// between tiling levels.
-static IREE::Codegen::LoweringConfigAttrInterface getNewLoweringConfig(
-    IREE::Codegen::LoweringConfigAttrInterface origLoweringConfig,
-    ArrayRef<IREE::CPU::LoweringConfigLevelInfo> tilingInfo,
-    bool setDistributionConfig) {
-  assert((isa<IREE::Codegen::LoweringConfigAttr, IREE::CPU::LoweringConfigAttr>(
-      origLoweringConfig)));
-  MLIRContext *ctx = origLoweringConfig.getContext();
-  if (isa<IREE::Codegen::LoweringConfigAttr>(origLoweringConfig)) {
-    TileSizesListType tileSizesList;
-    ScalableTileFlagsListType scalableTileFlagsList;
-    for (auto [level, tileSizes, scalableFlags] : tilingInfo) {
-      (void)level;
-      tileSizesList.push_back(tileSizes);
-      scalableTileFlagsList.push_back(scalableFlags);
-    }
-    return IREE::Codegen::LoweringConfigAttr::get(ctx, tileSizesList,
-                                                  scalableTileFlagsList);
-  }
-
+/// scalable tile flags. The distribution tiling sizes is not set if it is
+/// false.
+static IREE::Codegen::LoweringConfigAttrInterface
+getNewLoweringConfig(MLIRContext *ctx,
+                     ArrayRef<IREE::CPU::LoweringConfigLevelInfo> tilingInfo,
+                     bool setDistributionConfig) {
   SmallVector<NamedAttribute> newItems;
   for (auto [level, tileSizes, scalableFlags] : tilingInfo) {
     if (!setDistributionConfig && level == TilingLevel::DistributionTiles) {
@@ -2751,13 +2733,16 @@ static LogicalResult
 adjustTileSizesForUnPackOp(mlir::FunctionOpInterface entryPointFn,
                            Operation *rootOp) {
   auto linalgOp = dyn_cast<linalg::LinalgOp>(rootOp);
-  if (!linalgOp)
+  if (!linalgOp) {
+    return success();
+  }
+  auto loweringConfig =
+      getLoweringConfig<IREE::CPU::LoweringConfigAttr>(linalgOp);
+  if (!loweringConfig) {
+    // Tile size adjustment is only available when the rootOp uses
+    // IREE::CPU::LoweringConfigAttr.
     return success();
-  IREE::Codegen::LoweringConfigAttrInterface loweringConfig =
-      getLoweringConfig(linalgOp);
-  std::unique_ptr<TilingConfig> tilingConfig =
-      TilingConfig::create(loweringConfig);
-  TileSizesListType tileSizesList = tilingConfig->getTileSizes();
+  }
 
   bool foundUnPackOp = false;
   SmallVector<int64_t> alignedSizes(linalgOp.getNumLoops(), 1);
@@ -2792,7 +2777,7 @@ adjustTileSizesForUnPackOp(mlir::FunctionOpInterface entryPointFn,
 
   // Fixup for making tileSizes be multiple of inner_tile_sizes.
   SmallVector<IREE::CPU::LoweringConfigLevelInfo> tilingInfo =
-      tilingConfig->getTilingLevelInfo();
+      loweringConfig.getAvailableTilingInfo();
   for (IREE::CPU::LoweringConfigLevelInfo &info : tilingInfo) {
     SmallVector<int64_t> &tileSizes = info.sizes;
     for (auto idx : llvm::seq<int64_t>(0, tileSizes.size())) {
@@ -2824,7 +2809,7 @@ adjustTileSizesForUnPackOp(mlir::FunctionOpInterface entryPointFn,
   }
 
   IREE::Codegen::LoweringConfigAttrInterface newLoweringConfig =
-      getNewLoweringConfig(loweringConfig, tilingInfo,
+      getNewLoweringConfig(rootOp->getContext(), tilingInfo,
                            /*setDistributionConfig=*/true);
   return setOpConfigAndEntryPointFnTranslation(
       entryPointFn, rootOp, newLoweringConfig, pipeline, /*workgroupSize=*/{},
@@ -2947,7 +2932,13 @@ setLoweringConfigForComputeOps(mlir::FunctionOpInterface entryPointFn,
     return success();
   }
 
-  auto rootLoweringConfig = getLoweringConfig(rootOperation);
+  auto rootLoweringConfig =
+      getLoweringConfig<IREE::CPU::LoweringConfigAttr>(rootOperation);
+  if (!rootLoweringConfig) {
+    // Propagation is only available for IREE::CPU::LoweringConfigAttr.
+    return success();
+  }
+
   SmallVector<int64_t> distTileSizes, parallelVecTileSizes;
   SmallVector<bool> distScalableTileSizes, parallelVecScalableTileSizes;
   assert(rootLoweringConfig.hasWorkgroupTilingLevel());
@@ -3084,19 +3075,17 @@ 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;
     // For root op, we patch the adjusted tile sizes on its original tiling
     // config.
     if (op == rootOperation) {
-      newTilingInfo = tilingConfig->getTilingLevelInfo();
+      newTilingInfo = rootLoweringConfig.getAvailableTilingInfo();
       updateOrAddTilingLevelInfo(newTilingInfo, IREE::CPU::DistributionTiles,
                                  distTileSizes, distScalableTileSizes);
-      if (tilingConfig->getNumTilingLevels() > 1) {
+      if (rootLoweringConfig.hasTilingLevel(
+              IREE::CPU::VectorCommonParallelTiles)) {
         updateOrAddTilingLevelInfo(
             newTilingInfo, IREE::CPU::VectorCommonParallelTiles,
             commonVecTileSizes, commonVecScalableTileFlags);
@@ -3110,27 +3099,30 @@ setLoweringConfigForComputeOps(mlir::FunctionOpInterface entryPointFn,
                                  distTileSizes, falseVec);
       // The cache level tiling sizes are not adjusted, so we use the
       // config from the rootOp directly.
-      if (tilingConfig->isValidLevel(IREE::CPU::CacheParallelTiles)) {
-        updateOrAddTilingLevelInfo(newTilingInfo, IREE::CPU::CacheParallelTiles,
-                                   tilingConfig->getCacheParallelSizes(),
-                                   falseVec);
+      if (rootLoweringConfig.hasTilingLevel(IREE::CPU::CacheParallelTiles)) {
+        updateOrAddTilingLevelInfo(
+            newTilingInfo, IREE::CPU::CacheParallelTiles,
+            rootLoweringConfig.getStaticTilingLevelSizes(
+                IREE::CPU::CacheParallelTiles, rootOperation),
+            falseVec);
       }
-      if (tilingConfig->isValidLevel(IREE::CPU::CacheReductionTiles)) {
+      if (rootLoweringConfig.hasTilingLevel(IREE::CPU::CacheReductionTiles)) {
         updateOrAddTilingLevelInfo(
             newTilingInfo, IREE::CPU::CacheReductionTiles,
-            tilingConfig->getCacheReductionSizes(), falseVec);
+            rootLoweringConfig.getStaticTilingLevelSizes(
+                IREE::CPU::CacheReductionTiles, rootOperation),
+            falseVec);
       }
       updateOrAddTilingLevelInfo(
           newTilingInfo, IREE::CPU::VectorCommonParallelTiles,
           commonVecTileSizes, commonVecScalableTileFlags);
       bool setUpOK =
           TypeSwitch<Operation *, bool>(op)
               .Case<linalg::PackOp>([&](auto packOp) {
-                for (ArrayRef<bool> flags :
-                     tilingConfig->getScalableTileFlags()) {
-                  // TODO: Handle scalable flags
-                  if (llvm::any_of(flags, [&](bool flag) { return flag; }))
-                    return false;
+                // TODO: Handle scalable flags
+                if (llvm::any_of(rootLoweringConfig.getVectorScalableFlags(),
+                                 [&](bool flag) { return flag; })) {
+                  return false;
                 }
                 updateOrAddTilingLevelInfo(newTilingInfo,
                                            IREE::CPU::VectorReductionTiles,
@@ -3202,7 +3194,7 @@ setLoweringConfigForComputeOps(mlir::FunctionOpInterface entryPointFn,
                 return lhs.level < rhs.level;
               });
     IREE::Codegen::LoweringConfigAttrInterface config =
-        getNewLoweringConfig(rootLoweringConfig, newTilingInfo,
+        getNewLoweringConfig(rootOperation->getContext(), newTilingInfo,
                              /*setDistributionConfig=*/op == rootOperation);
     setLoweringConfig(op, config);
   }

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

@@ -1917,7 +1917,7 @@ func.func @custom_op(%arg0 : tensor<384x512xf32>, %arg1 : tensor<512x128xf32>,
   } -> tensor<384x128xf32>
   return %1 : tensor<384x128xf32>
 }
-//  CHECK-DAG: #[[CONFIG0:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[48, 64]]>
+//  CHECK-DAG: #[[CONFIG0:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[48, 64, 0]]>
 //  CHECK-DAG: #[[CONFIG1:.+]] = #iree_cpu.lowering_config<cache_parallel = [48, 64], vector_common_parallel = [8, 32]>
 //  CHECK-DAG: #[[CONFIG2:.+]] = #iree_cpu.lowering_config<cache_parallel = [48, 64, 0], distribution = [48, 64, 0], vector_common_parallel = [8, 32, 0], vector_reduction = [0, 0, 16]>
 //  CHECK-DAG: #[[TRANSLATION_INFO:.+]] = #iree_codegen.translation_info<pipeline = CPUDoubleTilingExpert, {enable_loop_peeling}>