[CPU] Take native_vector_size into accounts for attention op tiling. (#17349)
The second tile size maps the matmul M/N dimensions after decomposition.
The revision makes it have more reasonable tile sizes for the matmuls.
We are not able to force the size to native_vector_size because of the
bufferization issue, see https://github.com/iree-org/iree/issues/16956
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
index f5c047c..8802ec6 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
@@ -1692,10 +1692,10 @@
SmallVector<int64_t> distTileSizes = getDefaultDistributedLevelTileSizes(
attnOp, DistributionHeuristicConfig{});
int64_t iterationDomainRank = attnOp.getIterationDomainRank();
- // There are some dimensions are not tiled. Set vector tile sizes being ones
- // to avoid huge vectors.
- // TODO: We should be able to tile other dimensions.
SmallVector<int64_t> vecTileSizes(iterationDomainRank, 1);
+ int64_t vectorSize = getVectorSize(entryPointFn, attnOp.getOutputType());
+ vecTileSizes.back() = getMaxVectorTileSize(
+ /*numElem=*/distTileSizes.back(), vectorSize, vectorSize);
TileSizesListType tileSizes = {distTileSizes, vecTileSizes};
return setOpConfigAndEntryPointFnTranslation(
entryPointFn, attnOp, tileSizes,
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 96a4a01..5241b02 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
@@ -1613,7 +1613,7 @@
return
}
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[20, 64], [1, 1]]>
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[20, 64], [1, 32]]>
// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<CPULinalgExtTileAndVectorize>
// CHECK: func.func @attention()
// CHECK-SAME: translation_info = #[[TRANSLATION]]
