[Flow] Allow element-wise fusion of multi-reduction ops (#16503) This enables element-wise fusion of ops with multiple reduction dimensions in FusionOfTensorOps.
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/FusionOfTensorOps.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/FusionOfTensorOps.cpp index ea4f55a..060439e 100644 --- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/FusionOfTensorOps.cpp +++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/FusionOfTensorOps.cpp
@@ -19,6 +19,7 @@ #include "llvm/Support/Debug.h" #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h" #include "mlir/Dialect/Linalg/Transforms/Transforms.h" #include "mlir/Dialect/Math/IR/Math.h" #include "mlir/Dialect/MemRef/Transforms/Transforms.h" @@ -61,7 +62,8 @@ } /// Check if the producer generic op is fusable with the consumer generic op. -static bool areFusableOps(MLIRContext *context, OpOperand *fusedOperand) { +static bool areFusableOps(MLIRContext *context, OpOperand *fusedOperand, + bool fuseMultiReduction) { Operation *producerOp = fusedOperand->get().getDefiningOp(); Operation *consumerOp = fusedOperand->getOwner(); if (!producerOp) @@ -115,11 +117,17 @@ linalgConsumerOp.getNumLoops()) { return true; } - if (linalgConsumerOp.getNumReductionLoops() != 1 || - !linalgConsumerOp.getMatchingIndexingMap(fusedOperand) + if (!linalgConsumerOp.getMatchingIndexingMap(fusedOperand) .isPermutation()) { return false; } + if (!fuseMultiReduction && linalgConsumerOp.getNumReductionLoops() != 1) { + return false; + } + if (linalg::isaContractionOpInterface(linalgConsumerOp) || + linalg::isaConvolutionOpInterface(linalgConsumerOp)) { + return false; + } return true; } @@ -304,13 +312,15 @@ registry.insert<affine::AffineDialect, arith::ArithDialect, linalg::LinalgDialect, math::MathDialect>(); } - FusionOfTensorOpsPass(bool fuseMultiUse, unsigned multiUseFusionIteration) { + FusionOfTensorOpsPass(bool fuseMultiUse, bool fuseMultiReduction, + unsigned multiUseFusionIteration) { this->fuseMultiUse = fuseMultiUse; + this->fuseMultiReduction = fuseMultiReduction; this->multiUseFusionIteration = multiUseFusionIteration; } FusionOfTensorOpsPass(const FusionOfTensorOpsPass &pass) - : FusionOfTensorOpsPass(pass.fuseMultiUse, pass.multiUseFusionIteration) { - } + : FusionOfTensorOpsPass(pass.fuseMultiUse, pass.fuseMultiReduction, + pass.multiUseFusionIteration) {} void runOnOperation() override { Operation *funcOp = getOperation(); @@ -345,7 +355,7 @@ if (operands.size() >= kIreeMaxOperandCount) return false; - return areFusableOps(context, fusedOperand); + return areFusableOps(context, fusedOperand, fuseMultiReduction); }; linalg::populateElementwiseOpsFusionPatterns(fusionPatterns, fuseElementwiseOpsControlFn); @@ -501,10 +511,10 @@ } // namespace std::unique_ptr<InterfacePass<mlir::FunctionOpInterface>> -createFusionOfTensorOpsPass(bool fuseMultiUse, +createFusionOfTensorOpsPass(bool fuseMultiUse, bool fuseMultiReduction, unsigned multiUseFusionIteration) { - return std::make_unique<FusionOfTensorOpsPass>(fuseMultiUse, - multiUseFusionIteration); + return std::make_unique<FusionOfTensorOpsPass>( + fuseMultiUse, fuseMultiReduction, multiUseFusionIteration); } } // namespace mlir::iree_compiler::IREE::Flow
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.cpp index 2cdc605..56edc9e 100644 --- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.cpp +++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
@@ -80,6 +80,11 @@ llvm::cl::desc("Fuse multi-use ops."), llvm::cl::init(false)); +static llvm::cl::opt<bool> clEnableElementWiseFuseMultiReduction( + "iree-flow-element-wise-fuse-multi-reduction", + llvm::cl::desc("Enable element-wise fusion of multi-reduction loop ops."), + llvm::cl::init(true)); + static llvm::cl::opt<bool> clDispatchGenerateWorkloadRegion( "iree-flow-dispatch-generate-workload-region", llvm::cl::desc("Generate the workload region."), llvm::cl::init(true)); @@ -135,8 +140,10 @@ .addPass(mlir::createCanonicalizerPass) .addPass(mlir::createCSEPass) // Elementwise fusion. - .addPass( - []() { return createFusionOfTensorOpsPass(clEnableFuseMultiUse); }) + .addPass([]() { + return createFusionOfTensorOpsPass( + clEnableFuseMultiUse, clEnableElementWiseFuseMultiReduction); + }) .addPredicatedPass(clDetensoring, [&]() { return mlir::createLinalgDetensorizePass(); }) .addPass(mlir::createCanonicalizerPass)
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.h b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.h index 2bfc4bc..60cf678 100644 --- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.h +++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.h
@@ -65,6 +65,7 @@ // Creates a pass to fuse Linalg operations on tensors. std::unique_ptr<InterfacePass<mlir::FunctionOpInterface>> createFusionOfTensorOpsPass(bool fuseMultiUse = false, + bool fuseMultiReduction = true, unsigned multiUseFusionIteration = 2); // Create a pass to initialize all empty tensors after dispatch formation to
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.td b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.td index 55041e1..743b19e 100644 --- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.td +++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.td
@@ -168,6 +168,8 @@ let options = [ Option<"fuseMultiUse", "fuse-multi-use", "bool", /*default=*/"false", "Fuse ops with multiuse">, + Option<"fuseMultiReduction", "fuse-multi-reduction", "bool", + /*default=*/"true", "Fuse ops that have multiple reduction iterators">, Option<"multiUseFusionIteration", "multi-use-fusion-iteration", "unsigned", /*default=*/"2", "Number of iterations to fuse multiuse ops"> ];
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/fusion_of_tensor_ops.mlir b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/fusion_of_tensor_ops.mlir index f3741b6..da68ca1 100644 --- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/fusion_of_tensor_ops.mlir +++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/fusion_of_tensor_ops.mlir
@@ -366,10 +366,9 @@ // ----- -util.func public @nofuse_by_expand_dequant(%arg0 : tensor<11008x4096xi4>, %arg1 : tensor<11008x32x1xf16>, %arg2 : tensor<11008x32x1xf16>) -> (tensor<11008xf16>) { +util.func public @nofuse_by_expand_dequant(%arg0 : tensor<11008x4096xi4>, %arg1 : tensor<11008x32x1xf16>, %arg2 : tensor<11008x32x1xf16>, %arg3 : tensor<1x1x32x128xf16>) -> (tensor<11008xf16>) { %cst_1 = arith.constant 0.000000e+00 : f16 %0 = tensor.empty() : tensor<11008x32x128xf16> - %1 = arith.constant dense<0.000000e+00> : tensor<1x1x32x128xf16> %expanded = tensor.expand_shape %arg0 [[0], [1, 2]] : tensor<11008x4096xi4> into tensor<11008x32x128xi4> %collapsed = tensor.collapse_shape %arg2 [[0], [1, 2]] : tensor<11008x32x1xf16> into tensor<11008x32xf16> %collapsed_2 = tensor.collapse_shape %arg1 [[0], [1, 2]] : tensor<11008x32x1xf16> into tensor<11008x32xf16> @@ -381,7 +380,7 @@ %15 = arith.mulf %14, %in_8 : f16 linalg.yield %15 : f16 } -> tensor<11008x32x128xf16> - %collapsed_3 = tensor.collapse_shape %1 [[0, 1, 2], [3]] : tensor<1x1x32x128xf16> into tensor<32x128xf16> + %collapsed_3 = tensor.collapse_shape %arg3 [[0, 1, 2], [3]] : tensor<1x1x32x128xf16> into tensor<32x128xf16> %3 = tensor.empty() : tensor<11008xf16> %4 = linalg.fill ins(%cst_1 : f16) outs(%3 : tensor<11008xf16>) -> tensor<11008xf16> %5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0)>], iterator_types = ["parallel", "reduction", "reduction"]} ins(%collapsed_3, %2 : tensor<32x128xf16>, tensor<11008x32x128xf16>) outs(%4 : tensor<11008xf16>) { @@ -399,7 +398,7 @@ // CHECK-NOT: tensor.collapse_shape %[[DEQUANT]] // CHECK: %[[MATVEC:.+]] = linalg.generic // CHECK-SAME: iterator_types = ["parallel", "reduction", "reduction"] -// CHECK-SAME: ins(%[[DEQUANT]] : tensor<11008x32x128xf16>) +// CHECK-SAME: ins({{.+}}%[[DEQUANT]] // -----
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/pipeline_tests.mlir b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/pipeline_tests.mlir index 5485f59..290d2d7 100644 --- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/pipeline_tests.mlir +++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/pipeline_tests.mlir
@@ -44,16 +44,11 @@ // CHECK: flow.executable private @[[EXECUTABLE0:[a-zA-Z0-9_]+]] // CHECK: func.func @[[FUNC0:[a-zA-Z0-9_x]+]] // CHECK: linalg.generic -// CHECK-SAME: ["parallel", "parallel"] -// CHECK: flow.executable private @[[EXECUTABLE1:[a-zA-Z0-9_]+]] -// CHECK: func.func @[[FUNC1:[a-zA-Z0-9_x]+]] -// CHECK: linalg.generic -// CHECK-SAME: ["reduction"] +// CHECK-SAME: ["reduction", "reduction"] +// CHECK-NOT: linalg.generic // CHECK: util.func public @main( // CHECK: %[[T0:.+]] = flow.dispatch @[[EXECUTABLE0]]::@[[FUNC0]] -// CHECK: %[[T1:.+]] = flow.tensor.reshape %[[T0]] : tensor<833x833xf32> -> tensor<693889xf32> -// CHECK: %[[T2:.+]] = flow.dispatch @[[EXECUTABLE1]]::@[[FUNC1]](%[[T1]]) -// CHECK: util.return %[[T2]] +// CHECK: util.return %[[T0]] // -----