1 /* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "llvm/ADT/STLExtras.h"
17 #include "mlir/IR/Attributes.h" // from @llvm-project
18 #include "mlir/IR/Builders.h" // from @llvm-project
19 #include "mlir/IR/BuiltinOps.h" // from @llvm-project
20 #include "mlir/IR/PatternMatch.h" // from @llvm-project
21 #include "mlir/Pass/Pass.h" // from @llvm-project
22 #include "mlir/Pass/PassManager.h" // from @llvm-project
23 #include "mlir/Pass/PassRegistry.h" // from @llvm-project
24 #include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
25 #include "mlir/Transforms/Passes.h" // from @llvm-project
26 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
27 #include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
28
29 #define DEBUG_TYPE "tf-gpu-op-fusion"
30
31 namespace mlir {
32 namespace TF {
33
34 namespace {
35
36 // GpuOpFusionPass is a pass performing fusion specific to GPU targets.
37 // This is an ad-hoc pass for now, but should be integrated with some notion
38 // of "target" in the MLIR pipeline in the future.
39 class GpuOpFusionPass : public PassWrapper<GpuOpFusionPass, FunctionPass> {
40 public:
41 void runOnFunction() final;
42 };
43
44 // %y:6 = "tf.FusedBatchNormV3"(%x, %scale, %offset, %mean, %variance)
45 // %0 = "tf.Relu"(%y#0)
46 // ->
47 // %y:6 = "tf._FusedBatchNormEx"(%x, %scale, %offset, %mean, %variance)
48 //
49 // Or:
50 // %y:6 = "tf.FusedBatchNormV3"(%x, %scale, %offset, %mean, %variance)
51 // %0 = "tf.AddV2"(%y#0, %side_input)
52 // %1 = "tf.Relu"(%0)
53 // ->
54 // %y:6 = "tf._FusedBatchNormEx"(%x, %scale, %offset, %mean, %variance,
55 // %side_input)
56 // TODO(aminim): we should revisit this as a declarative pattern.
57 // For the second pattern, there is not good way in the framework to handle the
58 // commutativity of the AddV2: we want the FusedBatchNormV3 on any side.
59 // Also we need some native calls to handle the "hasOneUse" aspects and the
60 // optional extra operands for the AddV2 case.
61 struct ReluToFusedBatchNorm : public OpRewritePattern<ReluOp> {
62 using OpRewritePattern<ReluOp>::OpRewritePattern;
63
matchAndRewritemlir::TF::__anonbcf1e4bc0111::ReluToFusedBatchNorm64 LogicalResult matchAndRewrite(ReluOp relu_op,
65 PatternRewriter &rewriter) const override {
66 Operation *relu_input = relu_op.features().getDefiningOp();
67 if (!relu_input) return failure();
68 auto batch_norm = dyn_cast_or_null<FusedBatchNormV3Op>(relu_input);
69 AddV2Op add_op;
70 Value side_input;
71 if (!batch_norm) {
72 // We don't have a FusedBatchNorm as input to the ReLu, but we can get
73 // through an AddV2 as well.
74 add_op = dyn_cast_or_null<AddV2Op>(relu_input);
75 if (!add_op) return failure();
76
77 batch_norm =
78 dyn_cast_or_null<FusedBatchNormV3Op>(add_op.x().getDefiningOp());
79 if (batch_norm) {
80 side_input = add_op.y();
81 } else {
82 // Didn't get a FusedBatchNorm on the LHS of the AddV2, try the RHS.
83 batch_norm =
84 dyn_cast_or_null<FusedBatchNormV3Op>(add_op.y().getDefiningOp());
85 if (!batch_norm) return failure();
86 side_input = add_op.x();
87 }
88 }
89 assert(batch_norm);
90 if (batch_norm.is_training()) return failure();
91 if (!batch_norm.y().hasOneUse()) return failure();
92
93 // Build the newly fused operation to replace the batch norm
94 OperationState state(batch_norm.getLoc(),
95 _FusedBatchNormExOp::getOperationName());
96 state.addOperands(batch_norm.getOperands());
97 if (side_input) state.operands.push_back(side_input);
98 state.addTypes(batch_norm.getResultTypes());
99 state.addAttributes(batch_norm.getAttrs());
100 Operation *op = rewriter.createOperation(state);
101 rewriter.replaceOp(batch_norm, op->getResults());
102
103 // Depending on the case, we may fuse the add, the relu, or both.
104 if (!add_op || add_op.z().hasOneUse()) {
105 // We fuse the Relu only if the add has a single use, otherwise we only
106 // fuse the add itself.
107 op->setAttr("activation_mode", rewriter.getStringAttr("Relu"));
108 rewriter.replaceOp(relu_op, op->getResult(0));
109 }
110 if (add_op) {
111 rewriter.replaceOp(add_op, op->getResult(0));
112 }
113
114 return success();
115 }
116 };
117
runOnFunction()118 void GpuOpFusionPass::runOnFunction() {
119 FuncOp func = getFunction();
120 OwningRewritePatternList patterns;
121 patterns.insert<ReluToFusedBatchNorm>(&getContext());
122 (void)applyPatternsAndFoldGreedily(func, std::move(patterns));
123 }
124
125 } // namespace
126
CreateGpuOpFusionPass()127 std::unique_ptr<OperationPass<FuncOp>> CreateGpuOpFusionPass() {
128 return std::make_unique<GpuOpFusionPass>();
129 }
130
131 static PassRegistration<GpuOpFusionPass> layout_assignment(
132 "tf-gpu-op-fusion", "Fusion optimization for GPU targets");
133
134 } // namespace TF
135 } // namespace mlir
136