[CPU] Remove special distribution tile sizes from setting matmul config. (#15968)
They were used for DoubleTilingPadExpert pipeline, and were specialized
for MiniLMBert models. The pipeline was retired and we have been
tracking other models today. We don't have to use the configs for all
defaults.
It does not regress MiniLMBert model. It saves 30 ms (25.35%↓) for
MobileBert_int8 and 4 seconds (10.70%↓) for Falcon7bGptqPT. It regresses
a little for DeepLabV3, but it improves other ~10 models.
The regression only happens on single-threaded, and I think they should
be fixed by cache-level tiling afterwards. We are likely sitting in a
sweet spot because we mixed the context with distribution config.
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
index ef70e26..b9c30da 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
@@ -1147,20 +1147,22 @@
auto [vecTileSizes, vecScalableFlags] =
getMatmulVectorSizes(entryPointFn, linalgOp, vectorSize, isQuantized);
+ auto vecPreProcStrategy = getVectorPreProcStrategy(linalgOp);
+ bool usePeelingPipeline =
+ vecPreProcStrategy == VectorPreProcStrategy::Peeling;
+
+ LLVM_DEBUG(KD_DBGS() << "Vector pre-processing strategy: "
+ << vecPreProcStrategy << "\n");
+
DistributionHeuristicConfig distConfig;
- auto targetAttr = IREE::HAL::ExecutableTargetAttr::lookup(entryPointFn);
-
- // Use the default distribution for the matmul loops.
- int64_t defaultMaxSize = clDefaultDistTileSize;
- if (isX86(targetAttr) || isRISCV(targetAttr) ||
- (isAArch64(targetAttr) && hasAnySVEFeature(targetAttr))) {
- defaultMaxSize = 128;
- }
-
+ distConfig.maxTileSizes.resize(numLoops, clDefaultDistTileSize);
+ distConfig.allowIncompleteTile =
+ vecPreProcStrategy != VectorPreProcStrategy::None;
+ distConfig.vectorSizeHints.resize(numLoops, vectorSize);
bool isBM = isa<linalg::BatchMatmulOp>(contractionOp.getOperation());
- distConfig.maxTileSizes.resize(numLoops, defaultMaxSize);
if (isBM) {
distConfig.maxTileSizes[0] = 1;
+ distConfig.vectorSizeHints[0] = 1;
}
// Compute cache-level tile sizes. Cache a dimension only if there are
@@ -1172,31 +1174,13 @@
// Choose the next non-zero tile size immediately after the distribution
// level to help compute the distribution tile sizes.
- SmallVector<int64_t> distTileSizes;
- auto vecPreProcStrategy = getVectorPreProcStrategy(linalgOp);
- bool usePeelingPipeline =
- vecPreProcStrategy == VectorPreProcStrategy::Peeling;
-
- LLVM_DEBUG(KD_DBGS() << "Vector pre-processing strategy: "
- << vecPreProcStrategy << "\n");
-
- if (usePeelingPipeline && isX86(targetAttr)) {
- // It's inspired from https://github.com/iree-org/iree-llvm-sandbox repo.
- // Sandbox has [[288, 128, 512], [12, 32, 1]] setup. We scale 288 to 192
- // because 288/12*8=192
- if (numLoops == 3) {
- distConfig.maxTileSizes[0] = 192;
- distConfig.maxTileSizes[1] = 128;
- }
+ for (auto [cacheTileSize, vecTileSize] :
+ llvm::zip_equal(cacheTileSizes, vecTileSizes)) {
+ int64_t minTileSize = cacheTileSize != 0 ? cacheTileSize : vecTileSize;
+ distConfig.minTileSizes.push_back(minTileSize);
}
-
- distConfig.minTileSizes = vecTileSizes;
- distConfig.allowIncompleteTile = true;
- distConfig.vectorSizeHints.resize(numLoops, vectorSize);
- if (isBM) {
- distConfig.vectorSizeHints[0] = 1;
- }
- distTileSizes = getDefaultDistributedLevelTileSizes(linalgOp, distConfig);
+ SmallVector<int64_t> distTileSizes =
+ getDefaultDistributedLevelTileSizes(linalgOp, distConfig);
// TODO: We set cache tile sizes to the distribution sizes for now (no-op) to
// make sure there are no performance changes. This will let us change the
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_lowering_strategy.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_lowering_strategy.mlir
index 2aed28a..843c2b6 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_lowering_strategy.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_lowering_strategy.mlir
@@ -100,7 +100,7 @@
}
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[128, 128, 0], [8, [16], 0], [0, 0, 1], [0, 0, 0]]>
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[64, 64, 0], [8, [16], 0], [0, 0, 1], [0, 0, 0]]>
// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<CPUDoubleTilingExpert>
// CHECK: hal.executable.export public @matmul_tensors
// CHECK-SAME: translation_info = #[[TRANSLATION]]
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_sve_lowering_strategy_peeling.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_sve_lowering_strategy_peeling.mlir
index 693407d..424ca3f 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_sve_lowering_strategy_peeling.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_sve_lowering_strategy_peeling.mlir
@@ -47,7 +47,7 @@
}
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[128, 128, 0], [128, 128, 0], [0, 0, 0], [8, [16], 0], [0, 0, 1], [0, 0, 0]]>
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[64, 64, 0], [64, 64, 0], [0, 0, 0], [8, [16], 0], [0, 0, 1], [0, 0, 0]]>
// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<CPUDoubleTilingPeelingExpert>
// CHECK: hal.executable.export public @matmul_tensors
// CHECK-SAME: translation_info = #[[TRANSLATION]]
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_riscv_lowering_strategy.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_riscv_lowering_strategy.mlir
index 7e48b40..fca8f1a 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_riscv_lowering_strategy.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_riscv_lowering_strategy.mlir
@@ -39,8 +39,8 @@
}
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[128, 64], [8, 32], [0, 0], [0, 0]]>
-// CHECK-DAG: #[[CONFIG2:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[128, 64, 0], [128, 64, 0], [0, 0, 0], [8, 32, 0], [0, 0, 1], [0, 0, 0]]>
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[64, 64], [8, 32], [0, 0], [0, 0]]>
+// CHECK-DAG: #[[CONFIG2:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[64, 64, 0], [64, 64, 0], [0, 0, 0], [8, 32, 0], [0, 0, 1], [0, 0, 0]]>
// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<CPUDoubleTilingPeelingExpert>
// CHECK: hal.executable.export public @matmul_riscv
// CHECK-SAME: translation_info = #[[TRANSLATION]]
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_x86_64_lowering_strategy.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_x86_64_lowering_strategy.mlir
index a66310f..05f0484 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_x86_64_lowering_strategy.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_x86_64_lowering_strategy.mlir
@@ -83,7 +83,7 @@
}
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[128, 0], [128, 0], [0, 0], [32, 0], [0, 16], [0, 0]]>
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[64, 0], [64, 0], [0, 0], [32, 0], [0, 16], [0, 0]]>
// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<CPUDoubleTilingPeelingExpert>
// CHECK: hal.executable.export public @matvec_dynamic
// CHECK-SAME: translation_info = #[[TRANSLATION]]
@@ -418,7 +418,7 @@
}
}
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[129, 8, 0], [129, 8, 0], [0, 0, 0], [8, 8, 0], [0, 0, 16], [0, 0, 0]]>
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[43, 8, 0], [43, 8, 0], [0, 0, 0], [8, 8, 0], [0, 0, 16], [0, 0, 0]]>
// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<CPUDoubleTilingPeelingExpert>
// CHECK: hal.executable.export public @matmul_partially_peel
// CHECK-SAME: translation_info = #[[TRANSLATION]]
@@ -977,7 +977,7 @@
}
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[192, 64, 0], [192, 64, 0], [0, 0, 0], [8, 32, 0], [0, 0, 16], [0, 0, 0]]>
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[64, 64, 0], [64, 64, 0], [0, 0, 0], [8, 32, 0], [0, 0, 16], [0, 0, 0]]>
// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<CPUDoubleTilingPeelingExpert>
// CHECK: hal.executable.export public @matmul_static
// CHECK-SAME: translation_info = #[[TRANSLATION]]
@@ -1079,7 +1079,7 @@
}
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[128, 128, 0], [128, 128, 0], [0, 0, 0], [8, 32, 0], [0, 0, 16], [0, 0, 0]]>
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[64, 64, 0], [64, 64, 0], [0, 0, 0], [8, 32, 0], [0, 0, 16], [0, 0, 0]]>
// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<CPUDoubleTilingPeelingExpert>
// CHECK: hal.executable.export public @matmul_i8_i8_i32_static
// CHECK-SAME: translation_info = #[[TRANSLATION]]
@@ -1132,7 +1132,7 @@
}
}
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[192, 0, 0], [192, 0, 0], [0, 0, 0], [8, 32, 0], [0, 0, 16], [0, 0, 0]]>
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[64, 0, 0], [64, 0, 0], [0, 0, 0], [8, 32, 0], [0, 0, 16], [0, 0, 0]]>
// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<CPUDoubleTilingPeelingExpert>
// CHECK: hal.executable.export public @gemm_unit_N
// CHECK-SAME: translation_info = #[[TRANSLATION]]
@@ -1225,7 +1225,7 @@
}
}
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[11, 32, 0], [11, 32, 0], [0, 0, 0], [8, 32, 0], [0, 0, 16], [0, 0, 0]]>
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[11, 49, 0], [11, 49, 0], [0, 0, 0], [8, 32, 0], [0, 0, 16], [0, 0, 0]]>
// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<CPUDoubleTilingPeelingExpert>
// CHECK: hal.executable.export public @matmul_odd
// CHECK-SAME: translation_info = #[[TRANSLATION]]
@@ -1984,7 +1984,7 @@
}
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[192, 144, 0], [192, 144, 0], [0, 0, 0], [8, 32, 0], [0, 0, 16], [0, 0, 0]]>
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[64, 48, 0], [64, 48, 0], [0, 0, 0], [8, 32, 0], [0, 0, 16], [0, 0, 0]]>
// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<CPUDoubleTilingPeelingExpert>
// CHECK: hal.executable.export public @quant_model
// CHECK-SAME: translation_info = #[[TRANSLATION]]
