[LLVMCPU] Make LLVMCPUTile prefer using the lowering_config of tiled op. (#13922)
If the target operation has its own lowering_config, use its own config
for tiling.
It is a step toward https://github.com/openxla/iree/issues/13706 and
https://github.com/openxla/iree/issues/13474
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTile.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTile.cpp
index c1a0d16..dc19e35 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTile.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTile.cpp
@@ -65,18 +65,12 @@
auto funcOp = getOperation();
SmallVector<Operation *> computeOps = getComputeOps(funcOp);
- FailureOr<IREE::Codegen::LoweringConfigAttr> maybeLoweringConfig =
+ FailureOr<IREE::Codegen::LoweringConfigAttr> rootLoweringConfig =
getLoweringConfig(computeOps);
- if (failed(maybeLoweringConfig)) {
+ if (failed(rootLoweringConfig)) {
LLVM_DEBUG(llvm::dbgs() << "can't find lowering_config, skip tiling\n");
return;
}
- SmallVector<int64_t> tileSizes =
- maybeLoweringConfig.value().getTileSizeVals(tilingLevel);
- if (llvm::all_of(tileSizes, [](int64_t v) { return v == 0; })) {
- LLVM_DEBUG(llvm::dbgs() << "tiling sizes are all zeros, skip tiling\n");
- return;
- }
for (auto computeOp : computeOps) {
auto op = cast<TilingInterface>(computeOp);
@@ -89,6 +83,17 @@
if (isa<tensor::PadOp>(computeOp)) continue;
LLVM_DEBUG(llvm::dbgs() << "candidate: " << op << "\n");
+ SmallVector<int64_t> tileSizes;
+ if (auto loweringConfig = getLoweringConfig(op)) {
+ tileSizes = loweringConfig.getTileSizeVals(tilingLevel);
+ } else {
+ tileSizes = rootLoweringConfig.value().getTileSizeVals(tilingLevel);
+ }
+
+ if (llvm::all_of(tileSizes, [](int64_t v) { return v == 0; })) {
+ LLVM_DEBUG(llvm::dbgs() << "tiling sizes are all zeros, skip tiling\n");
+ continue;
+ }
IRRewriter rewriter(context);
auto options = scf::SCFTilingOptions().setTileSizeComputationFunction(
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/BUILD.bazel b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/BUILD.bazel
index 87a2f94..8c605c0 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/BUILD.bazel
@@ -50,6 +50,7 @@
"synchronize_symbol_visibility.mlir",
"tensor_pad.mlir",
"test_config_mmt4d.mlir",
+ "tile.mlir",
"tile_and_fuse.mlir",
"transform_dialect_bufferize.mlir",
"transform_dialect_iree_tile_to_forall.mlir",
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/CMakeLists.txt
index 4ec3d6e..e76f445 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/CMakeLists.txt
@@ -45,6 +45,7 @@
"synchronize_symbol_visibility.mlir"
"tensor_pad.mlir"
"test_config_mmt4d.mlir"
+ "tile.mlir"
"tile_and_fuse.mlir"
"transform_dialect_bufferize.mlir"
"transform_dialect_iree_tile_to_forall.mlir"
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/tile.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/tile.mlir
new file mode 100644
index 0000000..fcfb3d4
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/tile.mlir
@@ -0,0 +1,41 @@
+// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-llvmcpu-tile{tiling-level=0}))" --split-input-file %s | FileCheck %s
+
+func.func @matmul_bias_add(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?xf32>) -> tensor<?x?xf32> {
+ %cst = arith.constant 0.0 : f32
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
+ %d1 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
+ %init = tensor.empty(%d0, %d1) : tensor<?x?xf32>
+ %0 = linalg.fill ins(%cst : f32) outs(%init : tensor<?x?xf32>) -> tensor<?x?xf32>
+ %1 = linalg.matmul {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[10, 20, 30]]>}
+ ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+ outs(%0 : tensor<?x?xf32>) -> tensor<?x?xf32>
+ %2 = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1)-> (d0, d1)>],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%1, %arg2 : tensor<?x?xf32>, tensor<?xf32>)
+ outs(%init : tensor<?x?xf32>) attrs = {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[1, 16]]>} {
+ ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
+ %3 = arith.addf %arg3, %arg4 : f32
+ linalg.yield %3 : f32
+ } -> tensor<?x?xf32>
+ return %2 : tensor<?x?xf32>
+}
+// CHECK-LABEL: func.func @matmul_bias_add(
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C10:.+]] = arith.constant 10 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[C20:.+]] = arith.constant 20 : index
+// CHECK-DAG: %[[C30:.+]] = arith.constant 30 : index
+// CHECK: scf.for {{.*}} step %[[C10]]
+// CHECK: scf.for {{.*}} step %[[C20]]
+// CHECK: linalg.fill
+// CHECK: scf.for {{.*}} step %[[C10]]
+// CHECK: scf.for {{.*}} step %[[C20]]
+// CHECK: scf.for {{.*}} step %[[C30]]
+// CHECK: linalg.matmul
+// CHECK: scf.for {{.*}} step %[[C1]]
+// CHECK: scf.for {{.*}} step %[[C16]]
+// CHECK: linalg.generic
