[LLVMCPU] Add multi lowering_config support for SplitReductionPass. (#14061)
It is a step toward https://github.com/openxla/iree/issues/13706
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUSplitReduction.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUSplitReduction.cpp
index 2179690..9632e3a 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUSplitReduction.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUSplitReduction.cpp
@@ -5,6 +5,7 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#include "iree/compiler/Codegen/LLVMCPU/LLVMCPUPasses.h"
+#include "iree/compiler/Codegen/LLVMCPU/TileSizeSelection.h"
#include "iree/compiler/Codegen/PassDetail.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
@@ -154,47 +155,45 @@
class LLVMCPUSplitReductionPass
: public LLVMCPUSplitReductionBase<LLVMCPUSplitReductionPass> {
public:
- LLVMCPUSplitReductionPass(bool fpReductionReordering)
- : fpReductionReordering(fpReductionReordering) {}
+ LLVMCPUSplitReductionPass(bool fpReductionReordering) {
+ this->enableFpReductionReordering = fpReductionReordering;
+ }
void getDependentDialects(DialectRegistry ®istry) const override {
registry.insert<linalg::LinalgDialect, scf::SCFDialect>();
}
void runOnOperation() override;
-
- private:
- bool fpReductionReordering = false;
};
void LLVMCPUSplitReductionPass::runOnOperation() {
MLIRContext *context = &getContext();
auto funcOp = getOperation();
- SmallVector<Operation *> computeOps = getComputeOps(funcOp);
- FailureOr<IREE::Codegen::LoweringConfigAttr> maybeLoweringConfig =
- getLoweringConfig(computeOps);
- if (failed(maybeLoweringConfig)) {
- LLVM_DEBUG(llvm::dbgs()
- << "can't find lowering_config, skip SplitReduction");
- return;
- }
- auto reductionSizes = maybeLoweringConfig.value().getTileSizeVals().back();
- if (reductionSizes.empty()) {
- LLVM_DEBUG(
- llvm::dbgs()
- << "the list of reduction tiling sizes is empty, skip SplitReduction");
- return;
- }
- int64_t size = reductionSizes.back();
-
IRRewriter rewriter(context);
SmallVector<linalg::GenericOp> candidates;
funcOp.walk([&](linalg::GenericOp op) { candidates.push_back(op); });
for (auto genericOp : candidates) {
LLVM_DEBUG(llvm::dbgs() << "candidate: " << genericOp << "\n");
- if (failed(splitReductionPrecondition(genericOp, fpReductionReordering))) {
+ if (failed(splitReductionPrecondition(genericOp,
+ enableFpReductionReordering))) {
continue;
}
+
+ FailureOr<IREE::Codegen::LoweringConfigAttr> maybeLoweringConfig =
+ getLoweringConfig(genericOp);
+ if (failed(maybeLoweringConfig)) {
+ LLVM_DEBUG(llvm::dbgs()
+ << "can't find lowering_config, skip SplitReduction");
+ continue;
+ }
+ TilingConfig tilingConfig(maybeLoweringConfig.value());
+ auto reductionSizes = tilingConfig.getVectorReductionSizes();
+ if (reductionSizes.empty()) {
+ LLVM_DEBUG(llvm::dbgs() << "the list of reduction tiling sizes is empty, "
+ "skip SplitReduction");
+ continue;
+ }
+ int64_t size = reductionSizes.back();
if (failed(splitReductionImpl(genericOp, size, rewriter))) {
return signalPassFailure();
}
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/BUILD.bazel b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/BUILD.bazel
index 8c605c0..f15a1b4 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/BUILD.bazel
@@ -46,6 +46,7 @@
"peel.mlir",
"peel_and_vectorize.mlir",
"pipeline_tests.mlir",
+ "split_reduction.mlir",
"split_reduction_pipeline_tests.mlir",
"synchronize_symbol_visibility.mlir",
"tensor_pad.mlir",
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/CMakeLists.txt
index e76f445..1ad9dae 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/CMakeLists.txt
@@ -41,6 +41,7 @@
"peel.mlir"
"peel_and_vectorize.mlir"
"pipeline_tests.mlir"
+ "split_reduction.mlir"
"split_reduction_pipeline_tests.mlir"
"synchronize_symbol_visibility.mlir"
"tensor_pad.mlir"
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/split_reduction.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/split_reduction.mlir
new file mode 100644
index 0000000..d3db5ca
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/split_reduction.mlir
@@ -0,0 +1,54 @@
+// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-llvmcpu-split-reduction{enable-fp-reduction-reordering},cse,canonicalize))" --split-input-file %s | FileCheck %s
+
+#config = #iree_codegen.lowering_config<tile_sizes = [[2, 5, 32, 0], [1, 1, 8, 0], [0, 0, 0, 8]]>
+#config1 = #iree_codegen.lowering_config<tile_sizes = [[2, 5, 32, 0], [1, 1, 8, 0], [0, 0, 0, 16]]>
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+module {
+ func.func @softmax(%arg0: tensor<2x5x4096x4096xf32>) -> tensor<2x5x4096x4096xf32> {
+ %c0 = arith.constant 0 : index
+ %0 = tensor.empty() : tensor<2x5x4096x4096xf32>
+ %1 = tensor.empty() : tensor<2x5x4096x4096xf32>
+ %2 = tensor.empty() : tensor<2x5x4096xf32>
+ %cst = arith.constant -1.000000e+30 : f32
+ %3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<2x5x4096xf32>) -> tensor<2x5x4096xf32>
+ %4 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "reduction"]} ins(%arg0 : tensor<2x5x4096x4096xf32>) outs(%3 : tensor<2x5x4096xf32>) attrs = {lowering_config = #config} {
+ ^bb0(%in: f32, %out: f32):
+ %9 = arith.maxf %in, %out : f32
+ linalg.yield %9 : f32
+ } -> tensor<2x5x4096xf32>
+ %5 = linalg.generic {indexing_maps = [#map, #map1, #map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0, %4 : tensor<2x5x4096x4096xf32>, tensor<2x5x4096xf32>) outs(%1 : tensor<2x5x4096x4096xf32>) {
+ ^bb0(%in: f32, %in_1: f32, %out: f32):
+ %9 = arith.subf %in, %in_1 : f32
+ %10 = math.exp %9 : f32
+ linalg.yield %10 : f32
+ } -> tensor<2x5x4096x4096xf32>
+ %cst_0 = arith.constant 0.000000e+00 : f32
+ %6 = linalg.fill ins(%cst_0 : f32) outs(%2 : tensor<2x5x4096xf32>) -> tensor<2x5x4096xf32>
+ %7 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "reduction"]} ins(%5 : tensor<2x5x4096x4096xf32>) outs(%6 : tensor<2x5x4096xf32>) attrs = {lowering_config = #config1} {
+ ^bb0(%in: f32, %out: f32):
+ %9 = arith.addf %in, %out : f32
+ linalg.yield %9 : f32
+ } -> tensor<2x5x4096xf32>
+ %8 = linalg.generic {indexing_maps = [#map, #map1, #map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%5, %7 : tensor<2x5x4096x4096xf32>, tensor<2x5x4096xf32>) outs(%1 : tensor<2x5x4096x4096xf32>) {
+ ^bb0(%in: f32, %in_1: f32, %out: f32):
+ %9 = arith.divf %in, %in_1 : f32
+ linalg.yield %9 : f32
+ } -> tensor<2x5x4096x4096xf32>
+ return %8 : tensor<2x5x4096x4096xf32>
+ }
+}
+
+// CHECK: func.func @softmax
+// CHECK: scf.for
+// CHECK: scf.for
+// CHECK: scf.for
+// CHECK: scf.for
+// CHECK: %[[RES1:.+]] = linalg.generic
+// CHECK: scf.yield %[[RES1]] : tensor<1x1x1x8xf32>
+// CHECK: scf.for
+// CHECK: scf.for
+// CHECK: scf.for
+// CHECK: scf.for
+// CHECK: %[[RES2:.+]] = linalg.generic
+// CHECK: scf.yield %[[RES2]] : tensor<1x1x1x16xf32>
